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