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