1use crate::znx_base::ZnxInfos;
  2use crate::{alloc_aligned, Backend, DataView, DataViewMut, Module, VecZnx, VecZnxToMut, VecZnxToRef, ZnxSliceSize, ZnxView, ZnxViewMut};
  3use rand::seq::SliceRandom;
  4use rand_core::RngCore;
  5use rand_distr::{Distribution, weighted::WeightedIndex};
  6use sampling::source::Source;
  7
  8pub struct ScalarZnx<D> {
  9    data: D,
 10    n: usize,
 11    cols: usize,
 12}
 13
 14impl<D> ZnxInfos for ScalarZnx<D> {
 15    fn cols(&self) -> usize {
 16        self.cols
 17    }
 18
 19    fn rows(&self) -> usize {
 20        1
 21    }
 22
 23    fn n(&self) -> usize {
 24        self.n
 25    }
 26
 27    fn size(&self) -> usize {
 28        1
 29    }
 30}
 31
 32impl<D> ZnxSliceSize for ScalarZnx<D> {
 33    fn sl(&self) -> usize {
 34        self.n()
 35    }
 36}
 37
 38impl<D> DataView for ScalarZnx<D> {
 39    type D = D;
 40    fn data(&self) -> &Self::D {
 41        &self.data
 42    }
 43}
 44
 45impl<D> DataViewMut for ScalarZnx<D> {
 46    fn data_mut(&mut self) -> &mut Self::D {
 47        &mut self.data
 48    }
 49}
 50
 51impl<D: AsRef<[u8]>> ZnxView for ScalarZnx<D> {
 52    type Scalar = i64;
 53}
 54
 55impl<D: AsMut<[u8]> + AsRef<[u8]>> ScalarZnx<D> {
 56    pub fn fill_ternary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 57        let choices: [i64; 3] = [-1, 0, 1];
 58        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 59        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 60        self.at_mut(col, 0)
 61            .iter_mut()
 62            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 63    }
 64
 65    pub fn fill_ternary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 66        assert!(hw <= self.n());
 67        self.at_mut(col, 0)[..hw]
 68            .iter_mut()
 69            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 70        self.at_mut(col, 0).shuffle(source);
 71    }
 72}
 73
 74impl<D: From<Vec<u8>>> ScalarZnx<D> {
 75    pub(crate) fn bytes_of<S: Sized>(n: usize, cols: usize) -> usize {
 76        n * cols * size_of::<S>()
 77    }
 78
 79    pub(crate) fn new<S: Sized>(n: usize, cols: usize) -> Self {
 80        let data = alloc_aligned::<u8>(Self::bytes_of::<S>(n, cols));
 81        Self {
 82            data: data.into(),
 83            n,
 84            cols,
 85        }
 86    }
 87
 88    pub(crate) fn new_from_bytes<S: Sized>(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
 89        let data: Vec<u8> = bytes.into();
 90        assert!(data.len() == Self::bytes_of::<S>(n, cols));
 91        Self {
 92            data: data.into(),
 93            n,
 94            cols,
 95        }
 96    }
 97}
 98
 99pub type ScalarZnxOwned = ScalarZnx<Vec<u8>>;
100
101pub(crate) fn bytes_of_scalar_znx<B: Backend>(module: &Module<B>, cols: usize) -> usize {
102    ScalarZnxOwned::bytes_of::<i64>(module.n(), cols)
103}
104
105pub trait ScalarZnxAlloc {
106    fn bytes_of_scalar_znx(&self, cols: usize) -> usize;
107    fn new_scalar_znx(&self, cols: usize) -> ScalarZnxOwned;
108    fn new_scalar_znx_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxOwned;
109}
110
111impl<B: Backend> ScalarZnxAlloc for Module<B> {
112    fn bytes_of_scalar_znx(&self, cols: usize) -> usize {
113        ScalarZnxOwned::bytes_of::<i64>(self.n(), cols)
114    }
115    fn new_scalar_znx(&self, cols: usize) -> ScalarZnxOwned {
116        ScalarZnxOwned::new::<i64>(self.n(), cols)
117    }
118    fn new_scalar_znx_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarZnxOwned {
119        ScalarZnxOwned::new_from_bytes::<i64>(self.n(), cols, bytes)
120    }
121}
122
123impl<D> ScalarZnx<D> {
124    pub(crate) fn from_data(data: D, n: usize, cols: usize) -> Self {
125        Self { data, n, cols }
126    }
127}
128
129pub trait ScalarZnxToRef {
130    fn to_ref(&self) -> ScalarZnx<&[u8]>;
131}
132
133pub trait ScalarZnxToMut {
134    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]>;
135}
136
137impl ScalarZnxToMut for ScalarZnx<Vec<u8>> {
138    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
139        ScalarZnx {
140            data: self.data.as_mut_slice(),
141            n: self.n,
142            cols: self.cols,
143        }
144    }
145}
146
147impl VecZnxToMut for ScalarZnx<Vec<u8>>{
148    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
149        VecZnx {
150            data: self.data.as_mut_slice(),
151            n: self.n,
152            cols: self.cols,
153            size: 1,
154        }
155    }
156}
157
158impl ScalarZnxToRef for ScalarZnx<Vec<u8>> {
159    fn to_ref(&self) -> ScalarZnx<&[u8]> {
160        ScalarZnx {
161            data: self.data.as_slice(),
162            n: self.n,
163            cols: self.cols,
164        }
165    }
166}
167
168impl VecZnxToRef for ScalarZnx<Vec<u8>>{
169    fn to_ref(&self) -> VecZnx<&[u8]> {
170        VecZnx {
171            data: self.data.as_slice(),
172            n: self.n,
173            cols: self.cols,
174            size: 1,
175        }
176    }
177}
178
179impl ScalarZnxToMut for ScalarZnx<&mut [u8]> {
180    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
181        ScalarZnx {
182            data: self.data,
183            n: self.n,
184            cols: self.cols,
185        }
186    }
187}
188
189impl VecZnxToMut for ScalarZnx<&mut [u8]> {
190    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
191        VecZnx {
192            data: self.data,
193            n: self.n,
194            cols: self.cols,
195            size: 1,
196        }
197    }
198}
199
200impl ScalarZnxToRef for ScalarZnx<&mut [u8]> {
201    fn to_ref(&self) -> ScalarZnx<&[u8]> {
202        ScalarZnx {
203            data: self.data,
204            n: self.n,
205            cols: self.cols,
206        }
207    }
208}
209
210impl VecZnxToRef for ScalarZnx<&mut [u8]> {
211    fn to_ref(&self) -> VecZnx<&[u8]> {
212        VecZnx {
213            data: self.data,
214            n: self.n,
215            cols: self.cols,
216            size: 1,
217        }
218    }
219}
220
221impl ScalarZnxToRef for ScalarZnx<&[u8]> {
222    fn to_ref(&self) -> ScalarZnx<&[u8]> {
223        ScalarZnx {
224            data: self.data,
225            n: self.n,
226            cols: self.cols,
227        }
228    }
229}
230
231impl VecZnxToRef for ScalarZnx<&[u8]> {
232    fn to_ref(&self) -> VecZnx<&[u8]> {
233        VecZnx {
234            data: self.data,
235            n: self.n,
236            cols: self.cols,
237            size: 1,
238        }
239    }
240}