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, Reset, 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 ScalarZnx<Vec<u8>> {
111    pub fn alloc_bytes(n: usize, cols: usize) -> usize {
112        n * cols * size_of::<i64>()
113    }
114
115    pub fn alloc(n: usize, cols: usize) -> Self {
116        let data: Vec<u8> = alloc_aligned::<u8>(Self::alloc_bytes(n, cols));
117        Self {
118            data: data.into(),
119            n,
120            cols,
121        }
122    }
123
124    pub fn from_bytes(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
125        let data: Vec<u8> = bytes.into();
126        assert!(data.len() == Self::alloc_bytes(n, cols));
127        Self {
128            data: data.into(),
129            n,
130            cols,
131        }
132    }
133}
134
135impl<D: DataMut> ZnxZero for ScalarZnx<D> {
136    fn zero(&mut self) {
137        self.raw_mut().fill(0)
138    }
139    fn zero_at(&mut self, i: usize, j: usize) {
140        self.at_mut(i, j).fill(0);
141    }
142}
143
144impl<D: DataMut> FillUniform for ScalarZnx<D> {
145    fn fill_uniform(&mut self, source: &mut Source) {
146        source.fill_bytes(self.data.as_mut());
147    }
148}
149
150impl<D: DataMut> Reset for ScalarZnx<D> {
151    fn reset(&mut self) {
152        self.zero();
153        self.n = 0;
154        self.cols = 0;
155    }
156}
157
158pub type ScalarZnxOwned = ScalarZnx<Vec<u8>>;
159
160impl<D: Data> ScalarZnx<D> {
161    pub(crate) fn from_data(data: D, n: usize, cols: usize) -> Self {
162        Self { data, n, cols }
163    }
164}
165
166pub trait ScalarZnxToRef {
167    fn to_ref(&self) -> ScalarZnx<&[u8]>;
168}
169
170impl<D: DataRef> ScalarZnxToRef for ScalarZnx<D> {
171    fn to_ref(&self) -> ScalarZnx<&[u8]> {
172        ScalarZnx {
173            data: self.data.as_ref(),
174            n: self.n,
175            cols: self.cols,
176        }
177    }
178}
179
180pub trait ScalarZnxToMut {
181    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]>;
182}
183
184impl<D: DataMut> ScalarZnxToMut for ScalarZnx<D> {
185    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
186        ScalarZnx {
187            data: self.data.as_mut(),
188            n: self.n,
189            cols: self.cols,
190        }
191    }
192}
193
194impl<D: DataRef> ScalarZnx<D> {
195    pub fn as_vec_znx(&self) -> VecZnx<&[u8]> {
196        VecZnx {
197            data: self.data.as_ref(),
198            n: self.n,
199            cols: self.cols,
200            size: 1,
201            max_size: 1,
202        }
203    }
204}
205
206impl<D: DataMut> ScalarZnx<D> {
207    pub fn as_vec_znx_mut(&mut self) -> VecZnx<&mut [u8]> {
208        VecZnx {
209            data: self.data.as_mut(),
210            n: self.n,
211            cols: self.cols,
212            size: 1,
213            max_size: 1,
214        }
215    }
216}
217
218use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
219
220impl<D: DataMut> ReaderFrom for ScalarZnx<D> {
221    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
222        self.n = reader.read_u64::<LittleEndian>()? as usize;
223        self.cols = reader.read_u64::<LittleEndian>()? as usize;
224        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
225        let buf: &mut [u8] = self.data.as_mut();
226        if buf.len() != len {
227            return Err(std::io::Error::new(
228                std::io::ErrorKind::UnexpectedEof,
229                format!("self.data.len()={} != read len={}", buf.len(), len),
230            ));
231        }
232        reader.read_exact(&mut buf[..len])?;
233        Ok(())
234    }
235}
236
237impl<D: DataRef> WriterTo for ScalarZnx<D> {
238    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
239        writer.write_u64::<LittleEndian>(self.n as u64)?;
240        writer.write_u64::<LittleEndian>(self.cols as u64)?;
241        let buf: &[u8] = self.data.as_ref();
242        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
243        writer.write_all(buf)?;
244        Ok(())
245    }
246}