1use rand::seq::SliceRandom;
  2use rand_core::RngCore;
  3use rand_distr::{Distribution, weighted::WeightedIndex};
  4use sampling::source::Source;
  5
  6use crate::{
  7    alloc_aligned,
  8    hal::{
  9        api::{DataView, DataViewMut, FillUniform, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero},
 10        layouts::{Data, DataMut, DataRef, ReaderFrom, VecZnx, WriterTo},
 11    },
 12};
 13
 14#[derive(PartialEq, Eq, Debug, Clone)]
 15pub struct ScalarZnx<D: Data> {
 16    pub(crate) data: D,
 17    pub(crate) n: usize,
 18    pub(crate) cols: usize,
 19}
 20
 21impl<D: Data> ZnxInfos for ScalarZnx<D> {
 22    fn cols(&self) -> usize {
 23        self.cols
 24    }
 25
 26    fn rows(&self) -> usize {
 27        1
 28    }
 29
 30    fn n(&self) -> usize {
 31        self.n
 32    }
 33
 34    fn size(&self) -> usize {
 35        1
 36    }
 37}
 38
 39impl<D: Data> ZnxSliceSize for ScalarZnx<D> {
 40    fn sl(&self) -> usize {
 41        self.n()
 42    }
 43}
 44
 45impl<D: Data> DataView for ScalarZnx<D> {
 46    type D = D;
 47    fn data(&self) -> &Self::D {
 48        &self.data
 49    }
 50}
 51
 52impl<D: Data> DataViewMut for ScalarZnx<D> {
 53    fn data_mut(&mut self) -> &mut Self::D {
 54        &mut self.data
 55    }
 56}
 57
 58impl<D: DataRef> ZnxView for ScalarZnx<D> {
 59    type Scalar = i64;
 60}
 61
 62impl<D: DataMut> ScalarZnx<D> {
 63    pub fn fill_ternary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 64        let choices: [i64; 3] = [-1, 0, 1];
 65        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 66        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 67        self.at_mut(col, 0)
 68            .iter_mut()
 69            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 70    }
 71
 72    pub fn fill_ternary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 73        assert!(hw <= self.n());
 74        self.at_mut(col, 0)[..hw]
 75            .iter_mut()
 76            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 77        self.at_mut(col, 0).shuffle(source);
 78    }
 79
 80    pub fn fill_binary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 81        let choices: [i64; 2] = [0, 1];
 82        let weights: [f64; 2] = [1.0 - prob, prob];
 83        let dist: WeightedIndex<f64> = WeightedIndex::new(&weights).unwrap();
 84        self.at_mut(col, 0)
 85            .iter_mut()
 86            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 87    }
 88
 89    pub fn fill_binary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 90        assert!(hw <= self.n());
 91        self.at_mut(col, 0)[..hw]
 92            .iter_mut()
 93            .for_each(|x: &mut i64| *x = (source.next_u32() & 1) as i64);
 94        self.at_mut(col, 0).shuffle(source);
 95    }
 96
 97    pub fn fill_binary_block(&mut self, col: usize, block_size: usize, source: &mut Source) {
 98        assert!(self.n() % block_size == 0);
 99        let max_idx: u64 = (block_size + 1) as u64;
100        let mask_idx: u64 = (1 << ((u64::BITS - max_idx.leading_zeros()) as u64)) - 1;
101        for block in self.at_mut(col, 0).chunks_mut(block_size) {
102            let idx: usize = source.next_u64n(max_idx, mask_idx) as usize;
103            if idx != block_size {
104                block[idx] = 1;
105            }
106        }
107    }
108}
109
110impl<D: DataRef> ScalarZnx<D> {
111    pub fn bytes_of(n: usize, cols: usize) -> usize {
112        n * cols * size_of::<i64>()
113    }
114}
115
116impl<D: DataRef + From<Vec<u8>>> ScalarZnx<D> {
117    pub fn alloc(n: usize, cols: usize) -> Self {
118        let data: Vec<u8> = alloc_aligned::<u8>(Self::bytes_of(n, cols));
119        Self {
120            data: data.into(),
121            n,
122            cols,
123        }
124    }
125
126    pub(crate) fn from_bytes(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
127        let data: Vec<u8> = bytes.into();
128        assert!(data.len() == Self::bytes_of(n, cols));
129        Self {
130            data: data.into(),
131            n,
132            cols,
133        }
134    }
135}
136
137impl<D: DataMut> ZnxZero for ScalarZnx<D> {
138    fn zero(&mut self) {
139        self.raw_mut().fill(0)
140    }
141    fn zero_at(&mut self, i: usize, j: usize) {
142        self.at_mut(i, j).fill(0);
143    }
144}
145
146impl<D: DataMut> FillUniform for ScalarZnx<D> {
147    fn fill_uniform(&mut self, source: &mut Source) {
148        source.fill_bytes(self.data.as_mut());
149    }
150}
151
152pub type ScalarZnxOwned = ScalarZnx<Vec<u8>>;
153
154impl<D: Data> ScalarZnx<D> {
155    pub(crate) fn from_data(data: D, n: usize, cols: usize) -> Self {
156        Self { data, n, cols }
157    }
158}
159
160pub trait ScalarZnxToRef {
161    fn to_ref(&self) -> ScalarZnx<&[u8]>;
162}
163
164impl<D: DataRef> ScalarZnxToRef for ScalarZnx<D> {
165    fn to_ref(&self) -> ScalarZnx<&[u8]> {
166        ScalarZnx {
167            data: self.data.as_ref(),
168            n: self.n,
169            cols: self.cols,
170        }
171    }
172}
173
174pub trait ScalarZnxToMut {
175    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]>;
176}
177
178impl<D: DataMut> ScalarZnxToMut for ScalarZnx<D> {
179    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
180        ScalarZnx {
181            data: self.data.as_mut(),
182            n: self.n,
183            cols: self.cols,
184        }
185    }
186}
187
188impl<D: DataRef> ScalarZnx<D> {
189    pub fn as_vec_znx(&self) -> VecZnx<&[u8]> {
190        VecZnx {
191            data: self.data.as_ref(),
192            n: self.n,
193            cols: self.cols,
194            size: 1,
195            max_size: 1,
196        }
197    }
198}
199
200impl<D: DataMut> ScalarZnx<D> {
201    pub fn as_vec_znx_mut(&mut self) -> VecZnx<&mut [u8]> {
202        VecZnx {
203            data: self.data.as_mut(),
204            n: self.n,
205            cols: self.cols,
206            size: 1,
207            max_size: 1,
208        }
209    }
210}
211
212use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
213
214impl<D: DataMut> ReaderFrom for ScalarZnx<D> {
215    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
216        self.n = reader.read_u64::<LittleEndian>()? as usize;
217        self.cols = reader.read_u64::<LittleEndian>()? as usize;
218        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
219        let buf: &mut [u8] = self.data.as_mut();
220        if buf.len() != len {
221            return Err(std::io::Error::new(
222                std::io::ErrorKind::UnexpectedEof,
223                format!("self.data.len()={} != read len={}", buf.len(), len),
224            ));
225        }
226        reader.read_exact(&mut buf[..len])?;
227        Ok(())
228    }
229}
230
231impl<D: DataRef> WriterTo for ScalarZnx<D> {
232    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
233        writer.write_u64::<LittleEndian>(self.n as u64)?;
234        writer.write_u64::<LittleEndian>(self.cols as u64)?;
235        let buf: &[u8] = self.data.as_ref();
236        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
237        writer.write_all(buf)?;
238        Ok(())
239    }
240}