1use base2k::{ZnxInfos, Layout, Module, VecZnx, VecZnxOps, MatZnxDft, MatZnxDftOps};
  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(module: &Module, log_base2k: usize, log_q: usize, size: usize, bytes: &mut [u8]) -> Elem<VecZnx>;
 12    fn from_bytes_borrow(module: &Module, log_base2k: usize, log_q: usize, size: usize, bytes: &mut [u8]) -> Elem<VecZnx>;
 13    fn bytes_of(module: &Module, log_base2k: usize, log_q: usize, size: usize) -> usize;
 14    fn zero(&mut self);
 15}
 16
 17impl ElemVecZnx for Elem<VecZnx> {
 18    fn bytes_of(module: &Module, log_base2k: usize, log_q: usize, size: usize) -> usize {
 19        let cols = (log_q + log_base2k - 1) / log_base2k;
 20        module.n() * cols * size * 8
 21    }
 22
 23    fn from_bytes(module: &Module, log_base2k: usize, log_q: usize, size: usize, bytes: &mut [u8]) -> Elem<VecZnx> {
 24        assert!(size > 0);
 25        let n: usize = module.n();
 26        assert!(bytes.len() >= Self::bytes_of(module, log_base2k, log_q, size));
 27        let mut value: Vec<VecZnx> = Vec::new();
 28        let cols: usize = (log_q + log_base2k - 1) / log_base2k;
 29        let elem_size = VecZnx::bytes_of(n, size, cols);
 30        let mut ptr: usize = 0;
 31        (0..size).for_each(|_| {
 32            value.push(VecZnx::from_bytes(n, 1, cols, &mut bytes[ptr..]));
 33            ptr += elem_size
 34        });
 35        Self {
 36            value,
 37            log_q,
 38            log_base2k,
 39            log_scale: 0,
 40        }
 41    }
 42
 43    fn from_bytes_borrow(module: &Module, log_base2k: usize, log_q: usize, size: usize, bytes: &mut [u8]) -> 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 cols: usize = (log_q + log_base2k - 1) / log_base2k;
 49        let elem_size = VecZnx::bytes_of(n, 1, cols);
 50        let mut ptr: usize = 0;
 51        (0..size).for_each(|_| {
 52            value.push(VecZnx::from_bytes_borrow(n, 1, cols, &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 zero(&mut self) {
 64        self.value.iter_mut().for_each(|i| i.zero());
 65    }
 66}
 67
 68pub trait ElemCommon<T> {
 69    fn n(&self) -> usize;
 70    fn log_n(&self) -> usize;
 71    fn elem(&self) -> &Elem<T>;
 72    fn elem_mut(&mut self) -> &mut Elem<T>;
 73    fn size(&self) -> usize;
 74    fn layout(&self) -> Layout;
 75    fn rows(&self) -> usize;
 76    fn cols(&self) -> usize;
 77    fn log_base2k(&self) -> usize;
 78    fn log_q(&self) -> usize;
 79    fn log_scale(&self) -> usize;
 80    fn at(&self, i: usize) -> &T;
 81    fn at_mut(&mut self, i: usize) -> &mut T;
 82}
 83
 84impl<T: ZnxInfos> ElemCommon<T> for Elem<T> {
 85    fn n(&self) -> usize {
 86        self.value[0].n()
 87    }
 88
 89    fn log_n(&self) -> usize {
 90        self.value[0].log_n()
 91    }
 92
 93    fn elem(&self) -> &Elem<T> {
 94        self
 95    }
 96
 97    fn elem_mut(&mut self) -> &mut Elem<T> {
 98        self
 99    }
100
101    fn size(&self) -> usize {
102        self.value.len()
103    }
104
105    fn layout(&self) -> Layout {
106        self.value[0].layout()
107    }
108
109    fn rows(&self) -> usize {
110        self.value[0].rows()
111    }
112
113    fn cols(&self) -> usize {
114        self.value[0].cols()
115    }
116
117    fn log_base2k(&self) -> usize {
118        self.log_base2k
119    }
120
121    fn log_q(&self) -> usize {
122        self.log_q
123    }
124
125    fn log_scale(&self) -> usize {
126        self.log_scale
127    }
128
129    fn at(&self, i: usize) -> &T {
130        assert!(i < self.size());
131        &self.value[i]
132    }
133
134    fn at_mut(&mut self, i: usize) -> &mut T {
135        assert!(i < self.size());
136        &mut self.value[i]
137    }
138}
139
140impl Elem<VecZnx> {
141    pub fn new(module: &Module, log_base2k: usize, log_q: usize, rows: usize) -> Self {
142        assert!(rows > 0);
143        let cols: usize = (log_q + log_base2k - 1) / log_base2k;
144        let mut value: Vec<VecZnx> = Vec::new();
145        (0..rows).for_each(|_| value.push(module.new_vec_znx(1, cols)));
146        Self {
147            value,
148            log_q,
149            log_base2k,
150            log_scale: 0,
151        }
152    }
153}
154
155impl Elem<MatZnxDft> {
156    pub fn new(module: &Module, log_base2k: usize, size: usize, rows: usize, cols: usize) -> Self {
157        assert!(rows > 0);
158        assert!(cols > 0);
159        let mut value: Vec<MatZnxDft> = Vec::new();
160        (0..size).for_each(|_| value.push(module.new_vmp_pmat(1, rows, cols)));
161        Self {
162            value: value,
163            log_q: 0,
164            log_base2k: log_base2k,
165            log_scale: 0,
166        }
167    }
168}