1use crate::znx_base::ZnxInfos;
  2use crate::{Backend, DataView, DataViewMut, Module, ZnxView, ZnxViewMut, alloc_aligned};
  3use rand::seq::SliceRandom;
  4use rand_core::RngCore;
  5use rand_distr::{Distribution, weighted::WeightedIndex};
  6use sampling::source::Source;
  7
  8// pub const SCALAR_ZNX_ROWS: usize = 1;
  9// pub const SCALAR_ZNX_SIZE: usize = 1;
 10
 11pub struct Scalar<D> {
 12    data: D,
 13    n: usize,
 14    cols: usize,
 15}
 16
 17impl<D> ZnxInfos for Scalar<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    fn sl(&self) -> usize {
 35        self.n()
 36    }
 37}
 38
 39impl<D> DataView for Scalar<D> {
 40    type D = D;
 41    fn data(&self) -> &Self::D {
 42        &self.data
 43    }
 44}
 45
 46impl<D> DataViewMut for Scalar<D> {
 47    fn data_mut(&mut self) -> &mut Self::D {
 48        &mut self.data
 49    }
 50}
 51
 52impl<D: AsRef<[u8]>> ZnxView for Scalar<D> {
 53    type Scalar = i64;
 54}
 55
 56impl<D: AsMut<[u8]> + AsRef<[u8]>> Scalar<D> {
 57    pub fn fill_ternary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 58        let choices: [i64; 3] = [-1, 0, 1];
 59        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 60        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 61        self.at_mut(col, 0)
 62            .iter_mut()
 63            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 64    }
 65
 66    pub fn fill_ternary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 67        assert!(hw <= self.n());
 68        self.at_mut(col, 0)[..hw]
 69            .iter_mut()
 70            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 71        self.at_mut(col, 0).shuffle(source);
 72    }
 73
 74    // pub fn alias_as_vec_znx(&self) -> VecZnx {
 75    //     VecZnx {
 76    //         inner: ZnxBase {
 77    //             n: self.n(),
 78    //             rows: 1,
 79    //             cols: 1,
 80    //             size: 1,
 81    //             data: Vec::new(),
 82    //             ptr: self.ptr() as *mut u8,
 83    //         },
 84    //     }
 85    // }
 86}
 87
 88impl<D: From<Vec<u8>>> Scalar<D> {
 89    pub(crate) fn bytes_of<S: Sized>(n: usize, cols: usize) -> usize {
 90        n * cols * size_of::<S>()
 91    }
 92
 93    pub(crate) fn new<S: Sized>(n: usize, cols: usize) -> Self {
 94        let data = alloc_aligned::<u8>(Self::bytes_of::<S>(n, cols));
 95        Self {
 96            data: data.into(),
 97            n,
 98            cols,
 99        }
100    }
101
102    pub(crate) fn new_from_bytes<S: Sized>(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
103        let data: Vec<u8> = bytes.into();
104        assert!(data.len() == Self::bytes_of::<S>(n, cols));
105        Self {
106            data: data.into(),
107            n,
108            cols,
109        }
110    }
111}
112
113pub type ScalarOwned = Scalar<Vec<u8>>;
114
115pub trait ScalarAlloc {
116    fn bytes_of_scalar(&self, cols: usize) -> usize;
117    fn new_scalar(&self, cols: usize) -> ScalarOwned;
118    fn new_scalar_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarOwned;
119    // fn new_scalar_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> Scalar;
120}
121
122impl<B: Backend> ScalarAlloc for Module<B> {
123    fn bytes_of_scalar(&self, cols: usize) -> usize {
124        ScalarOwned::bytes_of::<i64>(self.n(), cols)
125    }
126    fn new_scalar(&self, cols: usize) -> ScalarOwned {
127        ScalarOwned::new::<i64>(self.n(), cols)
128    }
129    fn new_scalar_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> ScalarOwned {
130        ScalarOwned::new_from_bytes::<i64>(self.n(), cols, bytes)
131    }
132    // fn new_scalar_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> Scalar {
133    //     Scalar::from_bytes_borrow(self, SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE, bytes)
134    // }
135}
136
137// impl<B: Backend> ZnxAlloc<B> for Scalar {
138//     type Scalar = i64;
139
140//     fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, _size: usize, bytes: &mut [u8]) -> Self {
141//         Self {
142//             inner: ZnxBase::from_bytes_borrow(module.n(), SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE, bytes),
143//         }
144//     }
145
146//     fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, _size: usize) -> usize {
147//         debug_assert_eq!(
148//             _rows, SCALAR_ZNX_ROWS,
149//             "rows != {} not supported for Scalar",
150//             SCALAR_ZNX_ROWS
151//         );
152//         debug_assert_eq!(
153//             _size, SCALAR_ZNX_SIZE,
154//             "rows != {} not supported for Scalar",
155//             SCALAR_ZNX_SIZE
156//         );
157//         module.n() * cols * std::mem::size_of::<self::Scalar>()
158//     }
159// }