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