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