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 bool::*;
 27pub use decomposer::{Decomposer, DecomposerIter, DefaultDecomposer};
 28pub use ntt::{Ntt, NttBackendU64, NttInit};
 29pub use shortint::FheUint8;
 30
 31pub trait Matrix: AsRef<[Self::R]> {
 32    type MatElement;
 33    type R: Row<Element = Self::MatElement>;
 34
 35    fn dimension(&self) -> (usize, usize);
 36
 37    fn get_row(&self, row_idx: usize) -> impl Iterator<Item = &Self::MatElement> {
 38        self.as_ref()[row_idx].as_ref().iter().map(move |r| r)
 39    }
 40
 41    fn get_row_slice(&self, row_idx: usize) -> &[Self::MatElement] {
 42        self.as_ref()[row_idx].as_ref()
 43    }
 44
 45    fn iter_rows(&self) -> impl Iterator<Item = &Self::R> {
 46        self.as_ref().iter().map(move |r| r)
 47    }
 48
 49    fn get(&self, row_idx: usize, column_idx: usize) -> &Self::MatElement {
 50        &self.as_ref()[row_idx].as_ref()[column_idx]
 51    }
 52
 53    fn split_at_row(&self, idx: usize) -> (&[<Self as Matrix>::R], &[<Self as Matrix>::R]) {
 54        self.as_ref().split_at(idx)
 55    }
 56
 57    /// Does the matrix fit sub-matrix of dimension row x col
 58    fn fits(&self, row: usize, col: usize) -> bool;
 59}
 60
 61pub trait MatrixMut: Matrix + AsMut<[<Self as Matrix>::R]>
 62where
 63    <Self as Matrix>::R: RowMut,
 64{
 65    fn get_row_mut(&mut self, row_index: usize) -> &mut [Self::MatElement] {
 66        self.as_mut()[row_index].as_mut()
 67    }
 68
 69    fn iter_rows_mut(&mut self) -> impl Iterator<Item = &mut Self::R> {
 70        self.as_mut().iter_mut().map(move |r| r)
 71    }
 72
 73    fn set(&mut self, row_idx: usize, column_idx: usize, val: <Self as Matrix>::MatElement) {
 74        self.as_mut()[row_idx].as_mut()[column_idx] = val;
 75    }
 76
 77    fn split_at_row_mut(
 78        &mut self,
 79        idx: usize,
 80    ) -> (&mut [<Self as Matrix>::R], &mut [<Self as Matrix>::R]) {
 81        self.as_mut().split_at_mut(idx)
 82    }
 83}
 84
 85pub trait MatrixEntity: Matrix // where
 86// <Self as Matrix>::MatElement: Zero,
 87{
 88    fn zeros(row: usize, col: usize) -> Self;
 89}
 90
 91pub trait Row: AsRef<[Self::Element]> {
 92    type Element;
 93}
 94
 95pub trait RowMut: Row + AsMut<[<Self as Row>::Element]> {}
 96
 97pub trait RowEntity: Row {
 98    fn zeros(col: usize) -> Self;
 99}
100
101trait Secret {
102    type Element;
103    fn values(&self) -> &[Self::Element];
104}
105
106impl<T> Matrix for Vec<Vec<T>> {
107    type MatElement = T;
108    type R = Vec<T>;
109
110    fn dimension(&self) -> (usize, usize) {
111        (self.len(), self[0].len())
112    }
113
114    fn fits(&self, row: usize, col: usize) -> bool {
115        self.len() >= row && self[0].len() >= col
116    }
117}
118
119impl<T> Matrix for &[Vec<T>] {
120    type MatElement = T;
121    type R = Vec<T>;
122
123    fn dimension(&self) -> (usize, usize) {
124        (self.len(), self[0].len())
125    }
126
127    fn fits(&self, row: usize, col: usize) -> bool {
128        self.len() >= row && self[0].len() >= col
129    }
130}
131
132impl<T> Matrix for &mut [Vec<T>] {
133    type MatElement = T;
134    type R = Vec<T>;
135
136    fn dimension(&self) -> (usize, usize) {
137        (self.len(), self[0].len())
138    }
139
140    fn fits(&self, row: usize, col: usize) -> bool {
141        self.len() >= row && self[0].len() >= col
142    }
143}
144
145impl<T> MatrixMut for Vec<Vec<T>> {}
146impl<T> MatrixMut for &mut [Vec<T>] {}
147
148impl<T: Zero + Clone> MatrixEntity for Vec<Vec<T>> {
149    fn zeros(row: usize, col: usize) -> Self {
150        vec![vec![T::zero(); col]; row]
151    }
152}
153
154impl<T> Row for Vec<T> {
155    type Element = T;
156}
157
158impl<T> Row for [T] {
159    type Element = T;
160}
161
162impl<T> RowMut for Vec<T> {}
163
164impl<T: Zero + Clone> RowEntity for Vec<T> {
165    fn zeros(col: usize) -> Self {
166        vec![T::zero(); col]
167    }
168}
169
170pub trait Encryptor<M: ?Sized, C> {
171    fn encrypt(&self, m: &M) -> C;
172}
173
174pub trait Decryptor<M, C> {
175    fn decrypt(&self, c: &C) -> M;
176}
177
178pub trait MultiPartyDecryptor<M, C> {
179    type DecryptionShare;
180
181    fn gen_decryption_share(&self, c: &C) -> Self::DecryptionShare;
182    fn aggregate_decryption_shares(&self, c: &C, shares: &[Self::DecryptionShare]) -> M;
183}
184
185pub trait KeySwitchWithId<C> {
186    fn key_switch(&self, user_id: usize) -> C;
187}
188
189pub trait SampleExtractor<R> {
190    fn extract(&self, index: usize) -> R;
191}
192
193trait Encoder<F, T> {
194    fn encode(&self, v: F) -> T;
195}