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