1use base2k::{Infos, Module, VecZnx, 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: base2k::VecZnxCommon {}
 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
 78pub trait ElemCommon<T> {
 79    fn n(&self) -> usize;
 80    fn log_n(&self) -> usize;
 81    fn elem(&self) -> &Elem<T>;
 82    fn elem_mut(&mut self) -> &mut Elem<T>;
 83    fn size(&self) -> usize;
 84    fn rows(&self) -> usize;
 85    fn cols(&self) -> usize;
 86    fn log_base2k(&self) -> usize;
 87    fn log_q(&self) -> usize;
 88    fn log_scale(&self) -> usize;
 89    fn at(&self, i: usize) -> &T;
 90    fn at_mut(&mut self, i: usize) -> &mut T;
 91}
 92
 93impl<T: Infos> ElemCommon<T> for Elem<T> {
 94    fn n(&self) -> usize {
 95        self.value[0].n()
 96    }
 97
 98    fn log_n(&self) -> usize {
 99        self.value[0].log_n()
100    }
101
102    fn elem(&self) -> &Elem<T> {
103        self
104    }
105
106    fn elem_mut(&mut self) -> &mut Elem<T> {
107        self
108    }
109
110    fn size(&self) -> usize {
111        self.value.len()
112    }
113
114    fn rows(&self) -> usize {
115        self.value[0].rows()
116    }
117
118    fn cols(&self) -> usize {
119        self.value[0].cols()
120    }
121
122    fn log_base2k(&self) -> usize {
123        self.log_base2k
124    }
125
126    fn log_q(&self) -> usize {
127        self.log_q
128    }
129
130    fn log_scale(&self) -> usize {
131        self.log_scale
132    }
133
134    fn at(&self, i: usize) -> &T {
135        assert!(i < self.size());
136        &self.value[i]
137    }
138
139    fn at_mut(&mut self, i: usize) -> &mut T {
140        assert!(i < self.size());
141        &mut self.value[i]
142    }
143}
144
145impl Elem<VecZnx> {
146    pub fn new(module: &Module, log_base2k: usize, log_q: usize, rows: usize) -> Self {
147        assert!(rows > 0);
148        let limbs: usize = (log_q + log_base2k - 1) / log_base2k;
149        let mut value: Vec<VecZnx> = Vec::new();
150        (0..rows).for_each(|_| value.push(module.new_vec_znx(limbs)));
151        Self {
152            value,
153            log_q,
154            log_base2k,
155            log_scale: 0,
156        }
157    }
158}
159
160impl Elem<VmpPMat> {
161    pub fn new(module: &Module, log_base2k: usize, size: usize, rows: usize, cols: usize) -> Self {
162        assert!(rows > 0);
163        assert!(cols > 0);
164        let mut value: Vec<VmpPMat> = Vec::new();
165        (0..size).for_each(|_| value.push(module.new_vmp_pmat(rows, cols)));
166        Self {
167            value: value,
168            log_q: 0,
169            log_base2k: log_base2k,
170            log_scale: 0,
171        }
172    }
173}