1use base2k::{Infos, Module, VecZnx, VecZnxApi, VecZnxBorrow, VecZnxOps, VmpPMat, VmpPMatOps};
  2
  3use crate::parameters::Parameters;
  4
  5impl Parameters {
  6    pub fn elem_from_bytes<T>(&self, log_q: usize, size: usize, bytes: &mut [u8]) -> Elem<T>
  7    where
  8        T: VecZnxCommon<Owned = T>,
  9        Elem<T>: ElemVecZnx<T>,
 10    {
 11        Elem::<T>::from_bytes(self.module(), self.log_base2k(), log_q, size, bytes)
 12    }
 13}
 14
 15pub struct Elem<T> {
 16    pub value: Vec<T>,
 17    pub log_base2k: usize,
 18    pub log_q: usize,
 19    pub log_scale: usize,
 20}
 21
 22pub trait VecZnxCommon: VecZnxApi + Infos {}
 23impl VecZnxCommon for VecZnx {}
 24impl VecZnxCommon for VecZnxBorrow {}
 25
 26pub trait ElemVecZnx<T: VecZnxCommon<Owned = T>> {
 27    fn from_bytes(
 28        module: &Module,
 29        log_base2k: usize,
 30        log_q: usize,
 31        size: usize,
 32        bytes: &mut [u8],
 33    ) -> Elem<T>;
 34    fn bytes_of(module: &Module, log_base2k: usize, log_q: usize, size: usize) -> usize;
 35    fn zero(&mut self);
 36}
 37
 38impl<T> ElemVecZnx<T> for Elem<T>
 39where
 40    T: VecZnxCommon<Owned = T>,
 41{
 42    fn bytes_of(module: &Module, log_base2k: usize, log_q: usize, size: usize) -> usize {
 43        let cols = (log_q + log_base2k - 1) / log_base2k;
 44        module.n() * cols * size * 8
 45    }
 46
 47    fn from_bytes(
 48        module: &Module,
 49        log_base2k: usize,
 50        log_q: usize,
 51        size: usize,
 52        bytes: &mut [u8],
 53    ) -> Elem<T> {
 54        assert!(size > 0);
 55        let n: usize = module.n();
 56        assert!(bytes.len() >= Self::bytes_of(module, log_base2k, log_q, size));
 57        let mut value: Vec<T> = Vec::new();
 58        let limbs: usize = (log_q + log_base2k - 1) / log_base2k;
 59        let elem_size = T::bytes_of(n, limbs);
 60        let mut ptr: usize = 0;
 61        println!("{} {} {}", size, elem_size, bytes.len());
 62        (0..size).for_each(|_| {
 63            value.push(T::from_bytes(n, limbs, &mut bytes[ptr..]));
 64            ptr += elem_size
 65        });
 66        Self {
 67            value,
 68            log_q,
 69            log_base2k,
 70            log_scale: 0,
 71        }
 72    }
 73
 74    fn zero(&mut self) {
 75        self.value.iter_mut().for_each(|i| i.zero());
 76    }
 77}
 78
 79impl<T: Infos> Elem<T> {
 80    pub fn n(&self) -> usize {
 81        self.value[0].n()
 82    }
 83
 84    pub fn log_n(&self) -> usize {
 85        self.value[0].log_n()
 86    }
 87
 88    pub fn size(&self) -> usize {
 89        self.value.len()
 90    }
 91
 92    pub fn rows(&self) -> usize {
 93        self.value[0].rows()
 94    }
 95
 96    pub fn cols(&self) -> usize {
 97        self.value[0].cols()
 98    }
 99
100    pub fn log_base2k(&self) -> usize {
101        self.log_base2k
102    }
103
104    pub fn log_q(&self) -> usize {
105        self.log_q
106    }
107
108    pub fn log_scale(&self) -> usize {
109        self.log_scale
110    }
111
112    pub fn at(&self, i: usize) -> &T {
113        assert!(i < self.size());
114        &self.value[i]
115    }
116
117    pub fn at_mut(&mut self, i: usize) -> &mut T {
118        assert!(i < self.size());
119        &mut self.value[i]
120    }
121}
122
123impl Elem<VecZnx> {
124    pub fn new(module: &Module, log_base2k: usize, log_q: usize, rows: usize) -> Self {
125        assert!(rows > 0);
126        let limbs: usize = (log_q + log_base2k - 1) / log_base2k;
127        let mut value: Vec<VecZnx> = Vec::new();
128        (0..rows).for_each(|_| value.push(module.new_vec_znx(limbs)));
129        Self {
130            value,
131            log_q,
132            log_base2k,
133            log_scale: 0,
134        }
135    }
136}
137
138impl Elem<VmpPMat> {
139    pub fn new(module: &Module, log_base2k: usize, size: usize, rows: usize, cols: usize) -> Self {
140        assert!(rows > 0);
141        assert!(cols > 0);
142        let mut value: Vec<VmpPMat> = Vec::new();
143        (0..size).for_each(|_| value.push(module.new_vmp_pmat(rows, cols)));
144        Self {
145            value: value,
146            log_q: 0,
147            log_base2k: log_base2k,
148            log_scale: 0,
149        }
150    }
151}