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        (0..size).for_each(|_| {
 62            value.push(T::from_bytes(n, limbs, &mut bytes[ptr..]));
 63            ptr += elem_size
 64        });
 65        Self {
 66            value,
 67            log_q,
 68            log_base2k,
 69            log_scale: 0,
 70        }
 71    }
 72
 73    fn zero(&mut self) {
 74        self.value.iter_mut().for_each(|i| i.zero());
 75    }
 76}
 77
 78impl<T: Infos> Elem<T> {
 79    pub fn n(&self) -> usize {
 80        self.value[0].n()
 81    }
 82
 83    pub fn log_n(&self) -> usize {
 84        self.value[0].log_n()
 85    }
 86
 87    pub fn size(&self) -> usize {
 88        self.value.len()
 89    }
 90
 91    pub fn rows(&self) -> usize {
 92        self.value[0].rows()
 93    }
 94
 95    pub fn cols(&self) -> usize {
 96        self.value[0].cols()
 97    }
 98
 99    pub fn log_base2k(&self) -> usize {
100        self.log_base2k
101    }
102
103    pub fn log_q(&self) -> usize {
104        self.log_q
105    }
106
107    pub fn log_scale(&self) -> usize {
108        self.log_scale
109    }
110
111    pub fn at(&self, i: usize) -> &T {
112        assert!(i < self.size());
113        &self.value[i]
114    }
115
116    pub fn at_mut(&mut self, i: usize) -> &mut T {
117        assert!(i < self.size());
118        &mut self.value[i]
119    }
120}
121
122impl Elem<VecZnx> {
123    pub fn new(module: &Module, log_base2k: usize, log_q: usize, rows: usize) -> Self {
124        assert!(rows > 0);
125        let limbs: usize = (log_q + log_base2k - 1) / log_base2k;
126        let mut value: Vec<VecZnx> = Vec::new();
127        (0..rows).for_each(|_| value.push(module.new_vec_znx(limbs)));
128        Self {
129            value,
130            log_q,
131            log_base2k,
132            log_scale: 0,
133        }
134    }
135}
136
137impl Elem<VmpPMat> {
138    pub fn new(module: &Module, log_base2k: usize, size: usize, rows: usize, cols: usize) -> Self {
139        assert!(rows > 0);
140        assert!(cols > 0);
141        let mut value: Vec<VmpPMat> = Vec::new();
142        (0..size).for_each(|_| value.push(module.new_vmp_pmat(rows, cols)));
143        Self {
144            value: value,
145            log_q: 0,
146            log_base2k: log_base2k,
147            log_scale: 0,
148        }
149    }
150}