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    pub(crate) data: D,
 13    pub(crate) n: usize,
 14    pub(crate) 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    pub fn fill_binary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 77        let choices: [i64; 2] = [0, 1];
 78        let weights: [f64; 2] = [1.0 - prob, prob];
 79        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 80        self.at_mut(col, 0)
 81            .iter_mut()
 82            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 83    }
 84
 85    pub fn fill_binary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 86        assert!(hw <= self.n());
 87        self.at_mut(col, 0)[..hw]
 88            .iter_mut()
 89            .for_each(|x: &mut i64| *x = (source.next_u32() & 1) as i64);
 90        self.at_mut(col, 0).shuffle(source);
 91    }
 92
 93    pub fn fill_binary_block(&mut self, col: usize, block_size: usize, source: &mut Source) {
 94        assert!(self.n() % block_size == 0);
 95        for chunk in self.at_mut(col, 0).chunks_mut(block_size) {
 96            chunk[0] = 1;
 97            chunk.shuffle(source);
 98        }
 99    }
100}
101
102impl<D: From<Vec<u8>>> ScalarZnx<D> {
103    pub(crate) fn bytes_of<S: Sized>(n: usize, cols: usize) -> usize {
104        n * cols * size_of::<S>()
105    }
106
107    pub(crate) fn new<S: Sized>(n: usize, cols: usize) -> Self {
108        let data = alloc_aligned::<u8>(Self::bytes_of::<S>(n, cols));
109        Self {
110            data: data.into(),
111            n,
112            cols,
113        }
114    }
115
116    pub(crate) fn new_from_bytes<S: Sized>(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
117        let data: Vec<u8> = bytes.into();
118        assert!(data.len() == Self::bytes_of::<S>(n, cols));
119        Self {
120            data: data.into(),
121            n,
122            cols,
123        }
124    }
125}
126
127pub type ScalarZnxOwned = ScalarZnx<Vec<u8>>;
128
129pub(crate) fn bytes_of_scalar_znx<B: Backend>(module: &Module<B>, cols: usize) -> usize {
130    ScalarZnxOwned::bytes_of::<i64>(module.n(), cols)
131}
132
133pub trait ScalarZnxAlloc {
134    fn bytes_of_scalar_znx(&self, cols: usize) -> usize;
135    fn new_scalar_znx(&self, cols: usize) -> ScalarZnxOwned;
136    fn new_scalar_znx_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxOwned;
137}
138
139impl<B: Backend> ScalarZnxAlloc for Module<B> {
140    fn bytes_of_scalar_znx(&self, cols: usize) -> usize {
141        ScalarZnxOwned::bytes_of::<i64>(self.n(), cols)
142    }
143    fn new_scalar_znx(&self, cols: usize) -> ScalarZnxOwned {
144        ScalarZnxOwned::new::<i64>(self.n(), cols)
145    }
146    fn new_scalar_znx_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxOwned {
147        ScalarZnxOwned::new_from_bytes::<i64>(self.n(), cols, bytes)
148    }
149}
150
151pub trait ScalarZnxOps {
152    fn scalar_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
153    where
154        R: ScalarZnxToMut,
155        A: ScalarZnxToRef;
156
157    /// Applies the automorphism X^i -> X^ik on the selected column of `a`.
158    fn scalar_znx_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
159    where
160        A: ScalarZnxToMut;
161}
162
163impl<B: Backend> ScalarZnxOps for Module<B> {
164    fn scalar_znx_automorphism<R, A>(&self, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
165    where
166        R: ScalarZnxToMut,
167        A: ScalarZnxToRef,
168    {
169        let a: ScalarZnx<&[u8]> = a.to_ref();
170        let mut res: ScalarZnx<&mut [u8]> = res.to_mut();
171        #[cfg(debug_assertions)]
172        {
173            assert_eq!(a.n(), self.n());
174            assert_eq!(res.n(), self.n());
175        }
176        unsafe {
177            vec_znx::vec_znx_automorphism(
178                self.ptr,
179                k,
180                res.at_mut_ptr(res_col, 0),
181                res.size() as u64,
182                res.sl() as u64,
183                a.at_ptr(a_col, 0),
184                a.size() as u64,
185                a.sl() as u64,
186            )
187        }
188    }
189
190    fn scalar_znx_automorphism_inplace<A>(&self, k: i64, a: &mut A, a_col: usize)
191    where
192        A: ScalarZnxToMut,
193    {
194        let mut a: ScalarZnx<&mut [u8]> = a.to_mut();
195        #[cfg(debug_assertions)]
196        {
197            assert_eq!(a.n(), self.n());
198        }
199        unsafe {
200            vec_znx::vec_znx_automorphism(
201                self.ptr,
202                k,
203                a.at_mut_ptr(a_col, 0),
204                a.size() as u64,
205                a.sl() as u64,
206                a.at_ptr(a_col, 0),
207                a.size() as u64,
208                a.sl() as u64,
209            )
210        }
211    }
212}
213
214impl<D> ScalarZnx<D> {
215    pub(crate) fn from_data(data: D, n: usize, cols: usize) -> Self {
216        Self { data, n, cols }
217    }
218}
219
220pub trait ScalarZnxToRef {
221    fn to_ref(&self) -> ScalarZnx<&[u8]>;
222}
223
224impl<D> ScalarZnxToRef for ScalarZnx<D>
225where
226    D: AsRef<[u8]>,
227{
228    fn to_ref(&self) -> ScalarZnx<&[u8]> {
229        ScalarZnx {
230            data: self.data.as_ref(),
231            n: self.n,
232            cols: self.cols,
233        }
234    }
235}
236
237pub trait ScalarZnxToMut {
238    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]>;
239}
240
241impl<D> ScalarZnxToMut for ScalarZnx<D>
242where
243    D: AsRef<[u8]> + AsMut<[u8]>,
244{
245    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
246        ScalarZnx {
247            data: self.data.as_mut(),
248            n: self.n,
249            cols: self.cols,
250        }
251    }
252}
253
254impl<D> VecZnxToRef for ScalarZnx<D>
255where
256    D: AsRef<[u8]>,
257{
258    fn to_ref(&self) -> VecZnx<&[u8]> {
259        VecZnx {
260            data: self.data.as_ref(),
261            n: self.n,
262            cols: self.cols,
263            size: 1,
264        }
265    }
266}
267
268impl<D> VecZnxToMut for ScalarZnx<D>
269where
270    D: AsRef<[u8]> + AsMut<[u8]>,
271{
272    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
273        VecZnx {
274            data: self.data.as_mut(),
275            n: self.n,
276            cols: self.cols,
277            size: 1,
278        }
279    }
280}