1use crate::znx_base::{ZnxAlloc, ZnxBase, ZnxInfos, ZnxLayout, ZnxSliceSize};
  2use crate::{Backend, GetZnxBase, Module, VecZnx};
  3use rand::seq::SliceRandom;
  4use rand_core::RngCore;
  5use rand_distr::{Distribution, weighted::WeightedIndex};
  6use sampling::source::Source;
  7
  8pub const SCALAR_ZNX_ROWS: usize = 1;
  9pub const SCALAR_ZNX_SIZE: usize = 1;
 10
 11pub struct Scalar {
 12    pub inner: ZnxBase,
 13}
 14
 15impl GetZnxBase for Scalar {
 16    fn znx(&self) -> &ZnxBase {
 17        &self.inner
 18    }
 19
 20    fn znx_mut(&mut self) -> &mut ZnxBase {
 21        &mut self.inner
 22    }
 23}
 24
 25impl ZnxInfos for Scalar {}
 26
 27impl<B: Backend> ZnxAlloc<B> for Scalar {
 28    type Scalar = i64;
 29
 30    fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, _size: usize, bytes: &mut [u8]) -> Self {
 31        Self {
 32            inner: ZnxBase::from_bytes_borrow(module.n(), SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE, bytes),
 33        }
 34    }
 35
 36    fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, _size: usize) -> usize {
 37        debug_assert_eq!(
 38            _rows, SCALAR_ZNX_ROWS,
 39            "rows != {} not supported for Scalar",
 40            SCALAR_ZNX_ROWS
 41        );
 42        debug_assert_eq!(
 43            _size, SCALAR_ZNX_SIZE,
 44            "rows != {} not supported for Scalar",
 45            SCALAR_ZNX_SIZE
 46        );
 47        module.n() * cols * std::mem::size_of::<self::Scalar>()
 48    }
 49}
 50
 51impl ZnxLayout for Scalar {
 52    type Scalar = i64;
 53}
 54
 55impl ZnxSliceSize for Scalar {
 56    fn sl(&self) -> usize {
 57        self.n()
 58    }
 59}
 60
 61impl Scalar {
 62    pub fn fill_ternary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 63        let choices: [i64; 3] = [-1, 0, 1];
 64        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 65        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 66        self.at_mut(col, 0)
 67            .iter_mut()
 68            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 69    }
 70
 71    pub fn fill_ternary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 72        assert!(hw <= self.n());
 73        self.at_mut(col, 0)[..hw]
 74            .iter_mut()
 75            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 76        self.at_mut(col, 0).shuffle(source);
 77    }
 78
 79    pub fn alias_as_vec_znx(&self) -> VecZnx {
 80        VecZnx {
 81            inner: ZnxBase {
 82                n: self.n(),
 83                rows: 1,
 84                cols: 1,
 85                size: 1,
 86                data: Vec::new(),
 87                ptr: self.ptr() as *mut u8,
 88            },
 89        }
 90    }
 91}
 92
 93pub trait ScalarOps {
 94    fn bytes_of_scalar(&self, cols: usize) -> usize;
 95    fn new_scalar(&self, cols: usize) -> Scalar;
 96    fn new_scalar_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> Scalar;
 97    fn new_scalar_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> Scalar;
 98}
 99
100impl<B: Backend> ScalarOps for Module<B> {
101    fn bytes_of_scalar(&self, cols: usize) -> usize {
102        Scalar::bytes_of(self, SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE)
103    }
104    fn new_scalar(&self, cols: usize) -> Scalar {
105        Scalar::new(self, SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE)
106    }
107    fn new_scalar_from_bytes(&self, cols: usize, bytes: Vec<u8>) -> Scalar {
108        Scalar::from_bytes(self, SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE, bytes)
109    }
110    fn new_scalar_from_bytes_borrow(&self, cols: usize, bytes: &mut [u8]) -> Scalar {
111        Scalar::from_bytes_borrow(self, SCALAR_ZNX_ROWS, cols, SCALAR_ZNX_SIZE, bytes)
112    }
113}