1use itertools::izip;
  2use std::cmp::min;
  3
  4pub trait ZnxInfos {
  5    /// Returns the ring degree of the polynomials.
  6    fn n(&self) -> usize;
  7
  8    /// Returns the base two logarithm of the ring dimension of the polynomials.
  9    fn log_n(&self) -> usize {
 10        (usize::BITS - (self.n() - 1).leading_zeros()) as _
 11    }
 12
 13    /// Returns the number of rows.
 14    fn rows(&self) -> usize;
 15
 16    /// Returns the number of polynomials in each row.
 17    fn cols(&self) -> usize;
 18
 19    /// Returns the number of size per polynomial.
 20    fn size(&self) -> usize;
 21
 22    /// Returns the total number of small polynomials.
 23    fn poly_count(&self) -> usize {
 24        self.rows() * self.cols() * self.size()
 25    }
 26
 27    /// Returns the slice size, which is the offset between
 28    /// two size of the same column.
 29    fn sl(&self) -> usize;
 30}
 31
 32pub trait DataView {
 33    type D;
 34    fn data(&self) -> &Self::D;
 35}
 36
 37pub trait DataViewMut: DataView {
 38    fn data_mut(&mut self) -> &mut Self::D;
 39}
 40
 41pub trait ZnxView: ZnxInfos + DataView<D: AsRef<[u8]>> {
 42    type Scalar: Copy;
 43
 44    /// Returns a non-mutable pointer to the underlying coefficients array.
 45    fn as_ptr(&self) -> *const Self::Scalar {
 46        self.data().as_ref().as_ptr() as *const Self::Scalar
 47    }
 48
 49    /// Returns a non-mutable reference to the entire underlying coefficient array.
 50    fn raw(&self) -> &[Self::Scalar] {
 51        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
 52    }
 53
 54    /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
 55    fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
 56        #[cfg(debug_assertions)]
 57        {
 58            assert!(i < self.cols());
 59            assert!(j < self.size());
 60        }
 61        let offset: usize = self.n() * (j * self.cols() + i);
 62        unsafe { self.as_ptr().add(offset) }
 63    }
 64
 65    /// Returns non-mutable reference to the (i, j)-th small polynomial.
 66    fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
 67        unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
 68    }
 69}
 70
 71pub trait ZnxViewMut: ZnxView + DataViewMut<D: AsMut<[u8]>> {
 72    /// Returns a mutable pointer to the underlying coefficients array.
 73    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
 74        self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
 75    }
 76
 77    /// Returns a mutable reference to the entire underlying coefficient array.
 78    fn raw_mut(&mut self) -> &mut [Self::Scalar] {
 79        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
 80    }
 81
 82    /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
 83    fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
 84        #[cfg(debug_assertions)]
 85        {
 86            assert!(i < self.cols());
 87            assert!(j < self.size());
 88        }
 89        let offset: usize = self.n() * (j * self.cols() + i);
 90        unsafe { self.as_mut_ptr().add(offset) }
 91    }
 92
 93    /// Returns mutable reference to the (i, j)-th small polynomial.
 94    fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
 95        unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
 96    }
 97}
 98
 99//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
100impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: AsMut<[u8]>> {}
101
102pub trait ZnxZero: ZnxViewMut
103where
104    Self: Sized,
105{
106    fn zero(&mut self) {
107        unsafe {
108            std::ptr::write_bytes(
109                self.as_mut_ptr(),
110                0,
111                self.n() * size_of::<Self::Scalar>() * self.poly_count(),
112            );
113        }
114    }
115
116    fn zero_at(&mut self, i: usize, j: usize) {
117        unsafe {
118            std::ptr::write_bytes(
119                self.at_mut_ptr(i, j),
120                0,
121                self.n() * size_of::<Self::Scalar>(),
122            );
123        }
124    }
125}
126
127// Blanket implementations
128impl<T> ZnxZero for T where T: ZnxViewMut {}
129// impl<T> ZnxRsh for T where T: ZnxZero {}
130
131pub fn switch_degree<S: Copy, DMut: ZnxViewMut<Scalar = S> + ZnxZero, D: ZnxView<Scalar = S>>(
132    b: &mut DMut,
133    col_b: usize,
134    a: &D,
135    col_a: usize,
136) {
137    let (n_in, n_out) = (a.n(), b.n());
138    let (gap_in, gap_out): (usize, usize);
139
140    if n_in > n_out {
141        (gap_in, gap_out) = (n_in / n_out, 1)
142    } else {
143        (gap_in, gap_out) = (1, n_out / n_in);
144        b.zero();
145    }
146
147    let size: usize = min(a.size(), b.size());
148
149    (0..size).for_each(|i| {
150        izip!(
151            a.at(col_a, i).iter().step_by(gap_in),
152            b.at_mut(col_b, i).iter_mut().step_by(gap_out)
153        )
154        .for_each(|(x_in, x_out)| *x_out = *x_in);
155    });
156}
157
158use std::ops::{Add, AddAssign, Div, Mul, Neg, Shl, Shr, Sub};
159
160use crate::{ScratchBorr, cast_mut};
161pub trait Integer:
162    Copy
163    + Default
164    + PartialEq
165    + PartialOrd
166    + Add<Output = Self>
167    + Sub<Output = Self>
168    + Mul<Output = Self>
169    + Div<Output = Self>
170    + Neg<Output = Self>
171    + Shl<Output = Self>
172    + Shr<Output = Self>
173    + AddAssign
174{
175    const BITS: u32;
176}
177
178impl Integer for i64 {
179    const BITS: u32 = 64;
180}
181
182impl Integer for i128 {
183    const BITS: u32 = 128;
184}
185
186// (Jay)TODO: implement rsh for VecZnx, VecZnxBig
187// pub trait ZnxRsh: ZnxZero {
188//     fn rsh(&mut self, k: usize, log_base2k: usize, col: usize, carry: &mut [u8]) {
189//         rsh(k, log_base2k, self, col, carry)
190//     }
191// }
192pub fn rsh<V: ZnxZero>(k: usize, log_base2k: usize, a: &mut V, a_col: usize, scratch: &mut ScratchBorr)
193where
194    V::Scalar: From<usize> + Integer,
195{
196    let n: usize = a.n();
197    let size: usize = a.size();
198    let cols: usize = a.cols();
199
200    // #[cfg(debug_assertions)]
201    // {
202    //     assert!(
203    //         tmp_bytes.len() >= rsh_tmp_bytes::<V::Scalar>(n),
204    //         "invalid carry: carry.len()/size_ofSelf::Scalar={} < rsh_tmp_bytes({}, {})",
205    //         tmp_bytes.len() / size_of::<V::Scalar>(),
206    //         n,
207    //         size,
208    //     );
209    //     assert_alignement(tmp_bytes.as_ptr());
210    // }
211
212    let size: usize = a.size();
213    let steps: usize = k / log_base2k;
214
215    a.raw_mut().rotate_right(n * steps * cols);
216    (0..cols).for_each(|i| {
217        (0..steps).for_each(|j| {
218            a.zero_at(i, j);
219        })
220    });
221
222    let k_rem: usize = k % log_base2k;
223
224    if k_rem != 0 {
225        let (carry, _) = scratch.tmp_scalar_slice::<V::Scalar>(rsh_tmp_bytes::<V::Scalar>(n));
226
227        unsafe {
228            std::ptr::write_bytes(carry.as_mut_ptr(), 0, n * size_of::<V::Scalar>());
229        }
230
231        let log_base2k_t = V::Scalar::from(log_base2k);
232        let shift = V::Scalar::from(V::Scalar::BITS as usize - k_rem);
233        let k_rem_t = V::Scalar::from(k_rem);
234
235        (0..cols).for_each(|i| {
236            (steps..size).for_each(|j| {
237                izip!(carry.iter_mut(), a.at_mut(i, j).iter_mut()).for_each(|(ci, xi)| {
238                    *xi += *ci << log_base2k_t;
239                    *ci = (*xi << shift) >> shift;
240                    *xi = (*xi - *ci) >> k_rem_t;
241                });
242            });
243            //TODO:  ZERO CARRYcarry
244        })
245    }
246}
247
248pub fn rsh_tmp_bytes<T>(n: usize) -> usize {
249    n * std::mem::size_of::<T>()
250}