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