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, rows: usize, bytes: &mut [u8]) -> Elem<T>
  7    where
  8        T: VecZnxCommon,
  9        Elem<T>: Infos + ElemVecZnx<T>,
 10    {
 11        Elem::<T>::from_bytes(self.module(), self.log_base2k(), log_q, rows, 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> {
 27    fn from_bytes(
 28        module: &Module,
 29        log_base2k: usize,
 30        log_q: usize,
 31        rows: usize,
 32        bytes: &mut [u8],
 33    ) -> Elem<T>;
 34    fn bytes_of(module: &Module, log_base2k: usize, log_q: usize, rows: usize) -> usize;
 35    fn at(&self, i: usize) -> &T;
 36    fn at_mut(&mut self, i: usize) -> &mut T;
 37    fn zero(&mut self);
 38}
 39
 40impl<T> ElemVecZnx<T> for Elem<T>
 41where
 42    T: VecZnxCommon<Owned = T>,
 43    Elem<T>: Infos,
 44{
 45    fn bytes_of(module: &Module, log_base2k: usize, log_q: usize, rows: usize) -> usize {
 46        let cols = (log_q + log_base2k - 1) / log_base2k;
 47        module.n() * cols * (rows + 1) * 8
 48    }
 49
 50    fn from_bytes(
 51        module: &Module,
 52        log_base2k: usize,
 53        log_q: usize,
 54        rows: usize,
 55        bytes: &mut [u8],
 56    ) -> Elem<T> {
 57        assert!(rows > 0);
 58        let n: usize = module.n();
 59        assert!(bytes.len() >= Self::bytes_of(module, log_base2k, log_q, rows));
 60        let mut value: Vec<T> = Vec::new();
 61        let limbs: usize = (log_q + log_base2k - 1) / log_base2k;
 62        let size = T::bytes_of(n, limbs);
 63        let mut ptr: usize = 0;
 64        (0..rows).for_each(|_| {
 65            value.push(T::from_bytes(n, limbs, &mut bytes[ptr..]));
 66            ptr += size
 67        });
 68        Self {
 69            value,
 70            log_q,
 71            log_base2k,
 72            log_scale: 0,
 73        }
 74    }
 75
 76    fn at(&self, i: usize) -> &T {
 77        assert!(i < self.rows());
 78        &self.value[i]
 79    }
 80
 81    fn at_mut(&mut self, i: usize) -> &mut T {
 82        assert!(i < self.rows());
 83        &mut self.value[i]
 84    }
 85
 86    fn zero(&mut self) {
 87        self.value.iter_mut().for_each(|i| i.zero());
 88    }
 89}
 90
 91impl<T> Elem<T> {
 92    pub fn log_base2k(&self) -> usize {
 93        self.log_base2k
 94    }
 95
 96    pub fn log_q(&self) -> usize {
 97        self.log_q
 98    }
 99
100    pub fn log_scale(&self) -> usize {
101        self.log_scale
102    }
103}
104
105impl Infos for Elem<VecZnx> {
106    fn n(&self) -> usize {
107        self.value[0].n()
108    }
109
110    fn log_n(&self) -> usize {
111        self.value[0].log_n()
112    }
113
114    fn rows(&self) -> usize {
115        self.value.len()
116    }
117    fn cols(&self) -> usize {
118        self.value[0].cols()
119    }
120}
121
122impl Infos for Elem<VecZnxBorrow> {
123    fn n(&self) -> usize {
124        self.value[0].n()
125    }
126
127    fn log_n(&self) -> usize {
128        self.value[0].log_n()
129    }
130
131    fn rows(&self) -> usize {
132        self.value.len()
133    }
134    fn cols(&self) -> usize {
135        self.value[0].cols()
136    }
137}
138
139impl Elem<VecZnx> {
140    pub fn new(module: &Module, log_base2k: usize, log_q: usize, rows: usize) -> Self {
141        assert!(rows > 0);
142        let limbs: usize = (log_q + log_base2k - 1) / log_base2k;
143        let mut value: Vec<VecZnx> = Vec::new();
144        (0..rows).for_each(|_| value.push(module.new_vec_znx(limbs)));
145        Self {
146            value,
147            log_q,
148            log_base2k,
149            log_scale: 0,
150        }
151    }
152}
153
154impl Infos for Elem<VmpPMat> {
155    fn n(&self) -> usize {
156        self.value[0].n()
157    }
158
159    fn log_n(&self) -> usize {
160        self.value[0].log_n()
161    }
162
163    fn rows(&self) -> usize {
164        self.value[0].rows()
165    }
166
167    fn cols(&self) -> usize {
168        self.value[0].cols()
169    }
170}
171
172impl Elem<VmpPMat> {
173    pub fn new(module: &Module, log_base2k: usize, rows: usize, cols: usize) -> Self {
174        assert!(rows > 0);
175        assert!(cols > 0);
176        Self {
177            value: Vec::from([module.new_vmp_pmat(rows, cols); 1]),
178            log_q: 0,
179            log_base2k: log_base2k,
180            log_scale: 0,
181        }
182    }
183}