1use crate::ffi::vec_znx;
  2use crate::znx_base::ZnxInfos;
  3use crate::{
  4    Backend, DataView, DataViewMut, Module, VecZnx, VecZnxToMut, VecZnxToRef, ZnxSliceSize, ZnxView, ZnxViewMut, alloc_aligned,
  5};
  6use rand::seq::SliceRandom;
  7use rand_core::RngCore;
  8use rand_distr::{Distribution, weighted::WeightedIndex};
  9use sampling::source::Source;
 10
 11pub struct ScalarZnx<D> {
 12    data: D,
 13    n: usize,
 14    cols: usize,
 15}
 16
 17impl<D> ZnxInfos for ScalarZnx<D> {
 18    fn cols(&self) -> usize {
 19        self.cols
 20    }
 21
 22    fn rows(&self) -> usize {
 23        1
 24    }
 25
 26    fn n(&self) -> usize {
 27        self.n
 28    }
 29
 30    fn size(&self) -> usize {
 31        1
 32    }
 33}
 34
 35impl<D> ZnxSliceSize for ScalarZnx<D> {
 36    fn sl(&self) -> usize {
 37        self.n()
 38    }
 39}
 40
 41impl<D> DataView for ScalarZnx<D> {
 42    type D = D;
 43    fn data(&self) -> &Self::D {
 44        &self.data
 45    }
 46}
 47
 48impl<D> DataViewMut for ScalarZnx<D> {
 49    fn data_mut(&mut self) -> &mut Self::D {
 50        &mut self.data
 51    }
 52}
 53
 54impl<D: AsRef<[u8]>> ZnxView for ScalarZnx<D> {
 55    type Scalar = i64;
 56}
 57
 58impl<D: AsMut<[u8]> + AsRef<[u8]>> ScalarZnx<D> {
 59    pub fn fill_ternary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 60        let choices: [i64; 3] = [-1, 0, 1];
 61        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 62        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 63        self.at_mut(col, 0)
 64            .iter_mut()
 65            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 66    }
 67
 68    pub fn fill_ternary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 69        assert!(hw <= self.n());
 70        self.at_mut(col, 0)[..hw]
 71            .iter_mut()
 72            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 73        self.at_mut(col, 0).shuffle(source);
 74    }
 75}
 76
 77impl<D: From<Vec<u8>>> ScalarZnx<D> {
 78    pub(crate) fn bytes_of<S: Sized>(n: usize, cols: usize) -> usize {
 79        n * cols * size_of::<S>()
 80    }
 81
 82    pub(crate) fn new<S: Sized>(n: usize, cols: usize) -> Self {
 83        let data = alloc_aligned::<u8>(Self::bytes_of::<S>(n, cols));
 84        Self {
 85            data: data.into(),
 86            n,
 87            cols,
 88        }
 89    }
 90
 91    pub(crate) fn new_from_bytes<S: Sized>(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
 92        let data: Vec<u8> = bytes.into();
 93        assert!(data.len() == Self::bytes_of::<S>(n, cols));
 94        Self {
 95            data: data.into(),
 96            n,
 97            cols,
 98        }
 99    }
100}
101
102pub type ScalarZnxOwned = ScalarZnx<Vec<u8>>;
103
104pub(crate) fn bytes_of_scalar_znx<B: Backend>(module: &Module<B>, cols: usize) -> usize {
105    ScalarZnxOwned::bytes_of::<i64>(module.n(), cols)
106}
107
108pub trait ScalarZnxAlloc {
109    fn bytes_of_scalar_znx(&self, cols: usize) -> usize;
110    fn new_scalar_znx(&self, cols: usize) -> ScalarZnxOwned;
111    fn new_scalar_znx_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxOwned;
112}
113
114impl<B: Backend> ScalarZnxAlloc for Module<B> {
115    fn bytes_of_scalar_znx(&self, cols: usize) -> usize {
116        ScalarZnxOwned::bytes_of::<i64>(self.n(), cols)
117    }
118    fn new_scalar_znx(&self, cols: usize) -> ScalarZnxOwned {
119        ScalarZnxOwned::new::<i64>(self.n(), cols)
120    }
121    fn new_scalar_znx_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxOwned {
122        ScalarZnxOwned::new_from_bytes::<i64>(self.n(), cols, bytes)
123    }
124}
125
126pub trait ScalarZnxOps {
127    fn scalar_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
128    where
129        R: ScalarZnxToMut,
130        A: ScalarZnxToRef;
131
132    /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
133    fn vec_znx_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
134    where
135        A: ScalarZnxToMut;
136}
137
138impl<B: Backend> ScalarZnxOps for Module<B> {
139    fn scalar_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
140    where
141        R: ScalarZnxToMut,
142        A: ScalarZnxToRef,
143    {
144        let a: ScalarZnx<&[u8]> = a.to_ref();
145        let mut res: ScalarZnx<&mut [u8]> = res.to_mut();
146        #[cfg(debug_assertions)]
147        {
148            assert_eq!(a.n(), self.n());
149            assert_eq!(res.n(), self.n());
150        }
151        unsafe {
152            vec_znx::vec_znx_automorphism(
153                self.ptr,
154                k,
155                res.at_mut_ptr(res_col, 0),
156                res.size() as u64,
157                res.sl() as u64,
158                a.at_ptr(a_col, 0),
159                a.size() as u64,
160                a.sl() as u64,
161            )
162        }
163    }
164
165    fn vec_znx_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
166    where
167        A: ScalarZnxToMut,
168    {
169        let mut a: ScalarZnx<&mut [u8]> = a.to_mut();
170        #[cfg(debug_assertions)]
171        {
172            assert_eq!(a.n(), self.n());
173        }
174        unsafe {
175            vec_znx::vec_znx_automorphism(
176                self.ptr,
177                k,
178                a.at_mut_ptr(a_col, 0),
179                a.size() as u64,
180                a.sl() as u64,
181                a.at_ptr(a_col, 0),
182                a.size() as u64,
183                a.sl() as u64,
184            )
185        }
186    }
187}
188
189impl<D> ScalarZnx<D> {
190    pub(crate) fn from_data(data: D, n: usize, cols: usize) -> Self {
191        Self { data, n, cols }
192    }
193}
194
195pub trait ScalarZnxToRef {
196    fn to_ref(&self) -> ScalarZnx<&[u8]>;
197}
198
199pub trait ScalarZnxToMut {
200    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]>;
201}
202
203impl ScalarZnxToMut for ScalarZnx<Vec<u8>> {
204    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
205        ScalarZnx {
206            data: self.data.as_mut_slice(),
207            n: self.n,
208            cols: self.cols,
209        }
210    }
211}
212
213impl VecZnxToMut for ScalarZnx<Vec<u8>> {
214    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
215        VecZnx {
216            data: self.data.as_mut_slice(),
217            n: self.n,
218            cols: self.cols,
219            size: 1,
220        }
221    }
222}
223
224impl ScalarZnxToRef for ScalarZnx<Vec<u8>> {
225    fn to_ref(&self) -> ScalarZnx<&[u8]> {
226        ScalarZnx {
227            data: self.data.as_slice(),
228            n: self.n,
229            cols: self.cols,
230        }
231    }
232}
233
234impl VecZnxToRef for ScalarZnx<Vec<u8>> {
235    fn to_ref(&self) -> VecZnx<&[u8]> {
236        VecZnx {
237            data: self.data.as_slice(),
238            n: self.n,
239            cols: self.cols,
240            size: 1,
241        }
242    }
243}
244
245impl ScalarZnxToMut for ScalarZnx<&mut [u8]> {
246    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
247        ScalarZnx {
248            data: self.data,
249            n: self.n,
250            cols: self.cols,
251        }
252    }
253}
254
255impl VecZnxToMut for ScalarZnx<&mut [u8]> {
256    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
257        VecZnx {
258            data: self.data,
259            n: self.n,
260            cols: self.cols,
261            size: 1,
262        }
263    }
264}
265
266impl ScalarZnxToRef for ScalarZnx<&mut [u8]> {
267    fn to_ref(&self) -> ScalarZnx<&[u8]> {
268        ScalarZnx {
269            data: self.data,
270            n: self.n,
271            cols: self.cols,
272        }
273    }
274}
275
276impl VecZnxToRef for ScalarZnx<&mut [u8]> {
277    fn to_ref(&self) -> VecZnx<&[u8]> {
278        VecZnx {
279            data: self.data,
280            n: self.n,
281            cols: self.cols,
282            size: 1,
283        }
284    }
285}
286
287impl ScalarZnxToRef for ScalarZnx<&[u8]> {
288    fn to_ref(&self) -> ScalarZnx<&[u8]> {
289        ScalarZnx {
290            data: self.data,
291            n: self.n,
292            cols: self.cols,
293        }
294    }
295}
296
297impl VecZnxToRef for ScalarZnx<&[u8]> {
298    fn to_ref(&self) -> VecZnx<&[u8]> {
299        VecZnx {
300            data: self.data,
301            n: self.n,
302            cols: self.cols,
303            size: 1,
304        }
305    }
306}