1use crate::Module;
 2use rand::seq::SliceRandom;
 3use rand_core::RngCore;
 4use rand_distr::{Distribution, WeightedIndex};
 5use sampling::source::Source;
 6
 7pub struct Scalar(pub Vec<i64>);
 8
 9impl Module {
10    pub fn new_scalar(&self) -> Scalar {
11        Scalar::new(self.n())
12    }
13}
14
15impl Scalar {
16    pub fn new(n: usize) -> Self {
17        Self(vec![i64::default(); Self::buffer_size(n)])
18    }
19
20    pub fn buffer_size(n: usize) -> usize {
21        n
22    }
23
24    pub fn from_buffer(&mut self, n: usize, buf: &[i64]) {
25        let size = Self::buffer_size(n);
26        assert!(
27            buf.len() >= size,
28            "invalid buffer: buf.len()={} < self.buffer_size(n={})={}",
29            buf.len(),
30            n,
31            size
32        );
33        self.0 = Vec::from(&buf[..size])
34    }
35
36    pub fn as_ptr(&self) -> *const i64 {
37        self.0.as_ptr()
38    }
39
40    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
41        let choices: [i64; 3] = [-1, 0, 1];
42        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
43        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
44        self.0
45            .iter_mut()
46            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
47    }
48
49    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
50        self.0[..hw]
51            .iter_mut()
52            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
53        self.0.shuffle(source);
54    }
55}