poulpy/poulpy-hal/src/reference/znx/normalization.rs
poulpy-hal/src/reference/znx/normalization.rsRaw
1use itertools::izip;
2
3#[inline(always)]
4pub fn get_digit_i64(base2k: usize, x: i64) -> i64 {
5 (x << (u64::BITS - base2k as u32)) >> (u64::BITS - base2k as u32)
6}
7
8#[inline(always)]
9pub fn get_carry_i64(base2k: usize, x: i64, digit: i64) -> i64 {
10 (x.wrapping_sub(digit)) >> base2k
11}
12
13#[inline(always)]
14pub fn get_digit_i128(base2k: usize, x: i128) -> i128 {
15 (x << (u128::BITS - base2k as u32)) >> (u128::BITS - base2k as u32)
16}
17
18#[inline(always)]
19pub fn get_carry_i128(base2k: usize, x: i128, digit: i128) -> i128 {
20 (x.wrapping_sub(digit)) >> base2k
21}
22
23#[inline(always)]
24pub fn znx_normalize_first_step_carry_only_ref(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
25 #[cfg(debug_assertions)]
26 {
27 assert!(x.len() <= carry.len());
28 assert!(lsh < base2k);
29 }
30
31 if lsh == 0 {
32 x.iter().zip(carry.iter_mut()).for_each(|(x, c)| {
33 *c = get_carry_i64(base2k, *x, get_digit_i64(base2k, *x));
34 });
35 } else {
36 let base2k_lsh: usize = base2k - lsh;
37 x.iter().zip(carry.iter_mut()).for_each(|(x, c)| {
38 *c = get_carry_i64(base2k_lsh, *x, get_digit_i64(base2k_lsh, *x));
39 });
40 }
41}
42
43#[inline(always)]
44pub fn znx_normalize_first_step_inplace_ref(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
45 #[cfg(debug_assertions)]
46 {
47 assert!(x.len() <= carry.len());
48 assert!(lsh < base2k);
49 }
50
51 if lsh == 0 {
52 x.iter_mut().zip(carry.iter_mut()).for_each(|(x, c)| {
53 let digit: i64 = get_digit_i64(base2k, *x);
54 *c = get_carry_i64(base2k, *x, digit);
55 *x = digit;
56 });
57 } else {
58 let base2k_lsh: usize = base2k - lsh;
59 x.iter_mut().zip(carry.iter_mut()).for_each(|(x, c)| {
60 let digit: i64 = get_digit_i64(base2k_lsh, *x);
61 *c = get_carry_i64(base2k_lsh, *x, digit);
62 *x = digit << lsh;
63 });
64 }
65}
66
67#[inline(always)]
68pub fn znx_normalize_first_step_ref(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
69 #[cfg(debug_assertions)]
70 {
71 assert_eq!(x.len(), a.len());
72 assert!(x.len() <= carry.len());
73 assert!(lsh < base2k);
74 }
75
76 if lsh == 0 {
77 izip!(x.iter_mut(), a.iter(), carry.iter_mut()).for_each(|(x, a, c)| {
78 let digit: i64 = get_digit_i64(base2k, *a);
79 *c = get_carry_i64(base2k, *a, digit);
80 *x = digit;
81 });
82 } else {
83 let base2k_lsh: usize = base2k - lsh;
84 izip!(x.iter_mut(), a.iter(), carry.iter_mut()).for_each(|(x, a, c)| {
85 let digit: i64 = get_digit_i64(base2k_lsh, *a);
86 *c = get_carry_i64(base2k_lsh, *a, digit);
87 *x = digit << lsh;
88 });
89 }
90}
91
92#[inline(always)]
93pub fn znx_normalize_middle_step_carry_only_ref(base2k: usize, lsh: usize, x: &[i64], carry: &mut [i64]) {
94 #[cfg(debug_assertions)]
95 {
96 assert!(x.len() <= carry.len());
97 assert!(lsh < base2k);
98 }
99 if lsh == 0 {
100 x.iter().zip(carry.iter_mut()).for_each(|(x, c)| {
101 let digit: i64 = get_digit_i64(base2k, *x);
102 let carry: i64 = get_carry_i64(base2k, *x, digit);
103 let digit_plus_c: i64 = digit + *c;
104 *c = carry + get_carry_i64(base2k, digit_plus_c, get_digit_i64(base2k, digit_plus_c));
105 });
106 } else {
107 let base2k_lsh: usize = base2k - lsh;
108 x.iter().zip(carry.iter_mut()).for_each(|(x, c)| {
109 let digit: i64 = get_digit_i64(base2k_lsh, *x);
110 let carry: i64 = get_carry_i64(base2k_lsh, *x, digit);
111 let digit_plus_c: i64 = (digit << lsh) + *c;
112 *c = carry + get_carry_i64(base2k, digit_plus_c, get_digit_i64(base2k, digit_plus_c));
113 });
114 }
115}
116
117#[inline(always)]
118pub fn znx_normalize_middle_step_inplace_ref(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
119 #[cfg(debug_assertions)]
120 {
121 assert!(x.len() <= carry.len());
122 assert!(lsh < base2k);
123 }
124
125 if lsh == 0 {
126 x.iter_mut().zip(carry.iter_mut()).for_each(|(x, c)| {
127 let digit: i64 = get_digit_i64(base2k, *x);
128 let carry: i64 = get_carry_i64(base2k, *x, digit);
129 let digit_plus_c: i64 = digit + *c;
130 *x = get_digit_i64(base2k, digit_plus_c);
131 *c = carry + get_carry_i64(base2k, digit_plus_c, *x);
132 });
133 } else {
134 let base2k_lsh: usize = base2k - lsh;
135 x.iter_mut().zip(carry.iter_mut()).for_each(|(x, c)| {
136 let digit: i64 = get_digit_i64(base2k_lsh, *x);
137 let carry: i64 = get_carry_i64(base2k_lsh, *x, digit);
138 let digit_plus_c: i64 = (digit << lsh) + *c;
139 *x = get_digit_i64(base2k, digit_plus_c);
140 *c = carry + get_carry_i64(base2k, digit_plus_c, *x);
141 });
142 }
143}
144
145#[inline(always)]
146pub fn znx_extract_digit_addmul_ref(base2k: usize, lsh: usize, res: &mut [i64], src: &mut [i64]) {
147 for (r, s) in res.iter_mut().zip(src.iter_mut()) {
148 let digit: i64 = get_digit_i64(base2k, *s);
149 *s = get_carry_i64(base2k, *s, digit);
150 *r += digit << lsh;
151 }
152}
153
154#[inline(always)]
155pub fn znx_normalize_digit_ref(base2k: usize, res: &mut [i64], src: &mut [i64]) {
156 for (r, s) in res.iter_mut().zip(src.iter_mut()) {
157 let ri_digit: i64 = get_digit_i64(base2k, *r);
158 let ri_carry: i64 = get_carry_i64(base2k, *r, ri_digit);
159 *r = ri_digit;
160 *s += ri_carry;
161 }
162}
163
164#[inline(always)]
165pub fn znx_normalize_middle_step_ref(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
166 #[cfg(debug_assertions)]
167 {
168 assert_eq!(x.len(), a.len());
169 assert!(x.len() <= carry.len());
170 assert!(lsh < base2k);
171 }
172 if lsh == 0 {
173 izip!(x.iter_mut(), a.iter(), carry.iter_mut()).for_each(|(x, a, c)| {
174 let digit: i64 = get_digit_i64(base2k, *a);
175 let carry: i64 = get_carry_i64(base2k, *a, digit);
176 let digit_plus_c: i64 = digit + *c;
177 *x = get_digit_i64(base2k, digit_plus_c);
178 *c = carry + get_carry_i64(base2k, digit_plus_c, *x);
179 });
180 } else {
181 let base2k_lsh: usize = base2k - lsh;
182 izip!(x.iter_mut(), a.iter(), carry.iter_mut()).for_each(|(x, a, c)| {
183 let digit: i64 = get_digit_i64(base2k_lsh, *a);
184 let carry: i64 = get_carry_i64(base2k_lsh, *a, digit);
185 let digit_plus_c: i64 = (digit << lsh) + *c;
186 *x = get_digit_i64(base2k, digit_plus_c);
187 *c = carry + get_carry_i64(base2k, digit_plus_c, *x);
188 });
189 }
190}
191
192#[inline(always)]
193pub fn znx_normalize_final_step_inplace_ref(base2k: usize, lsh: usize, x: &mut [i64], carry: &mut [i64]) {
194 #[cfg(debug_assertions)]
195 {
196 assert!(x.len() <= carry.len());
197 assert!(lsh < base2k);
198 }
199
200 if lsh == 0 {
201 x.iter_mut().zip(carry.iter_mut()).for_each(|(x, c)| {
202 *x = get_digit_i64(base2k, get_digit_i64(base2k, *x) + *c);
203 });
204 } else {
205 let base2k_lsh: usize = base2k - lsh;
206 x.iter_mut().zip(carry.iter_mut()).for_each(|(x, c)| {
207 *x = get_digit_i64(base2k, (get_digit_i64(base2k_lsh, *x) << lsh) + *c);
208 });
209 }
210}
211
212#[inline(always)]
213pub fn znx_normalize_final_step_ref(base2k: usize, lsh: usize, x: &mut [i64], a: &[i64], carry: &mut [i64]) {
214 #[cfg(debug_assertions)]
215 {
216 assert!(x.len() <= carry.len());
217 assert!(lsh < base2k);
218 }
219 if lsh == 0 {
220 izip!(x.iter_mut(), a.iter(), carry.iter_mut()).for_each(|(x, a, c)| {
221 *x = get_digit_i64(base2k, get_digit_i64(base2k, *a) + *c);
222 });
223 } else {
224 let base2k_lsh: usize = base2k - lsh;
225 izip!(x.iter_mut(), a.iter(), carry.iter_mut()).for_each(|(x, a, c)| {
226 *x = get_digit_i64(base2k, (get_digit_i64(base2k_lsh, *a) << lsh) + *c);
227 });
228 }
229}