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
 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 scalar_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 scalar_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
199impl<D> ScalarZnxToRef for ScalarZnx<D>
200where
201    D: AsRef<[u8]>,
202{
203    fn to_ref(&self) -> ScalarZnx<&[u8]> {
204        ScalarZnx {
205            data: self.data.as_ref(),
206            n: self.n,
207            cols: self.cols,
208        }
209    }
210}
211
212pub trait ScalarZnxToMut {
213    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]>;
214}
215
216impl<D> ScalarZnxToMut for ScalarZnx<D>
217where
218    D: AsRef<[u8]> + AsMut<[u8]>,
219{
220    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
221        ScalarZnx {
222            data: self.data.as_mut(),
223            n: self.n,
224            cols: self.cols,
225        }
226    }
227}
228
229impl<D> VecZnxToRef for ScalarZnx<D>
230where
231    D: AsRef<[u8]>,
232{
233    fn to_ref(&self) -> VecZnx<&[u8]> {
234        VecZnx {
235            data: self.data.as_ref(),
236            n: self.n,
237            cols: self.cols,
238            size: 1,
239        }
240    }
241}
242
243impl<D> VecZnxToMut for ScalarZnx<D>
244where
245    D: AsRef<[u8]> + AsMut<[u8]>,
246{
247    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
248        VecZnx {
249            data: self.data.as_mut(),
250            n: self.n,
251            cols: self.cols,
252            size: 1,
253        }
254    }
255}