1use std::{iter::Once, sync::OnceLock};
  2
  3use itertools::{izip, Itertools};
  4use num_traits::{abs, Zero};
  5
  6mod backend;
  7mod bool;
  8mod decomposer;
  9mod lwe;
 10mod multi_party;
 11mod noise;
 12mod ntt;
 13mod pbs;
 14mod random;
 15mod rgsw;
 16mod shortint;
 17mod utils;
 18
 19pub use backend::{
 20    ArithmeticLazyOps, ArithmeticOps, ModInit, ModularOpsU64, ShoupMatrixFMA, VectorOps,
 21};
 22// pub use bool::{
 23//     aggregate_public_key_shares, aggregate_server_key_shares, gen_client_key,
 24// gen_mp_keys_phase1,     gen_mp_keys_phase2, set_mp_seed, set_parameter_set,
 25// ParameterSelector, };
 26pub use decomposer::{Decomposer, DecomposerIter, DefaultDecomposer};
 27pub use ntt::{Ntt, NttBackendU64, NttInit};
 28
 29pub trait Matrix: AsRef<[Self::R]> {
 30    type MatElement;
 31    type R: Row<Element = Self::MatElement>;
 32
 33    fn dimension(&self) -> (usize, usize);
 34
 35    fn get_row(&self, row_idx: usize) -> impl Iterator<Item = &Self::MatElement> {
 36        self.as_ref()[row_idx].as_ref().iter().map(move |r| r)
 37    }
 38
 39    fn get_row_slice(&self, row_idx: usize) -> &[Self::MatElement] {
 40        self.as_ref()[row_idx].as_ref()
 41    }
 42
 43    fn iter_rows(&self) -> impl Iterator<Item = &Self::R> {
 44        self.as_ref().iter().map(move |r| r)
 45    }
 46
 47    fn get(&self, row_idx: usize, column_idx: usize) -> &Self::MatElement {
 48        &self.as_ref()[row_idx].as_ref()[column_idx]
 49    }
 50
 51    fn split_at_row(&self, idx: usize) -> (&[<Self as Matrix>::R], &[<Self as Matrix>::R]) {
 52        self.as_ref().split_at(idx)
 53    }
 54
 55    /// Does the matrix fit sub-matrix of dimension row x col
 56    fn fits(&self, row: usize, col: usize) -> bool;
 57}
 58
 59pub trait MatrixMut: Matrix + AsMut<[<Self as Matrix>::R]>
 60where
 61    <Self as Matrix>::R: RowMut,
 62{
 63    fn get_row_mut(&mut self, row_index: usize) -> &mut [Self::MatElement] {
 64        self.as_mut()[row_index].as_mut()
 65    }
 66
 67    fn iter_rows_mut(&mut self) -> impl Iterator<Item = &mut Self::R> {
 68        self.as_mut().iter_mut().map(move |r| r)
 69    }
 70
 71    fn set(&mut self, row_idx: usize, column_idx: usize, val: <Self as Matrix>::MatElement) {
 72        self.as_mut()[row_idx].as_mut()[column_idx] = val;
 73    }
 74
 75    fn split_at_row_mut(
 76        &mut self,
 77        idx: usize,
 78    ) -> (&mut [<Self as Matrix>::R], &mut [<Self as Matrix>::R]) {
 79        self.as_mut().split_at_mut(idx)
 80    }
 81}
 82
 83pub trait MatrixEntity: Matrix // where
 84// <Self as Matrix>::MatElement: Zero,
 85{
 86    fn zeros(row: usize, col: usize) -> Self;
 87}
 88
 89pub trait Row: AsRef<[Self::Element]> {
 90    type Element;
 91}
 92
 93pub trait RowMut: Row + AsMut<[<Self as Row>::Element]> {}
 94
 95pub trait RowEntity: Row {
 96    fn zeros(col: usize) -> Self;
 97}
 98
 99trait Secret {
100    type Element;
101    fn values(&self) -> &[Self::Element];
102}
103
104impl<T> Matrix for Vec<Vec<T>> {
105    type MatElement = T;
106    type R = Vec<T>;
107
108    fn dimension(&self) -> (usize, usize) {
109        (self.len(), self[0].len())
110    }
111
112    fn fits(&self, row: usize, col: usize) -> bool {
113        self.len() >= row && self[0].len() >= col
114    }
115}
116
117impl<T> Matrix for &[Vec<T>] {
118    type MatElement = T;
119    type R = Vec<T>;
120
121    fn dimension(&self) -> (usize, usize) {
122        (self.len(), self[0].len())
123    }
124
125    fn fits(&self, row: usize, col: usize) -> bool {
126        self.len() >= row && self[0].len() >= col
127    }
128}
129
130impl<T> Matrix for &mut [Vec<T>] {
131    type MatElement = T;
132    type R = Vec<T>;
133
134    fn dimension(&self) -> (usize, usize) {
135        (self.len(), self[0].len())
136    }
137
138    fn fits(&self, row: usize, col: usize) -> bool {
139        self.len() >= row && self[0].len() >= col
140    }
141}
142
143impl<T> MatrixMut for Vec<Vec<T>> {}
144impl<T> MatrixMut for &mut [Vec<T>] {}
145
146impl<T: Zero + Clone> MatrixEntity for Vec<Vec<T>> {
147    fn zeros(row: usize, col: usize) -> Self {
148        vec![vec![T::zero(); col]; row]
149    }
150}
151
152impl<T> Row for Vec<T> {
153    type Element = T;
154}
155
156impl<T> Row for [T] {
157    type Element = T;
158}
159
160impl<T> RowMut for Vec<T> {}
161
162impl<T: Zero + Clone> RowEntity for Vec<T> {
163    fn zeros(col: usize) -> Self {
164        vec![T::zero(); col]
165    }
166}
167
168pub trait Encryptor<M: ?Sized, C> {
169    fn encrypt(&self, m: &M) -> C;
170}
171
172pub trait Decryptor<M, C> {
173    fn decrypt(&self, c: &C) -> M;
174}
175
176pub trait MultiPartyDecryptor<M, C> {
177    type DecryptionShare;
178
179    fn gen_decryption_share(&self, c: &C) -> Self::DecryptionShare;
180    fn aggregate_decryption_shares(&self, c: &C, shares: &[Self::DecryptionShare]) -> M;
181}
182
183pub trait KeySwitchWithId<C> {
184    fn key_switch(&self, user_id: usize) -> C;
185}