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