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