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