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