1use backend::{
  2    Backend, FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScratchOwned, VecZnxDft,
  3    ZnxInfos, ZnxZero,
  4};
  5use sampling::source::Source;
  6
  7use crate::{elem::Infos, glwe_ciphertext_fourier::GLWECiphertextFourier};
  8
  9#[derive(Clone, Copy, Debug)]
 10pub(crate) enum SecretDistribution {
 11    TernaryFixed(usize), // Ternary with fixed Hamming weight
 12    TernaryProb(f64),    // Ternary with probabilistic Hamming weight
 13    ZERO,                // Debug mod
 14    NONE,
 15}
 16
 17pub struct SecretKey<T> {
 18    pub(crate) data: ScalarZnx<T>,
 19    pub(crate) dist: SecretDistribution,
 20}
 21
 22impl SecretKey<Vec<u8>> {
 23    pub fn alloc<B: Backend>(module: &Module<B>, rank: usize) -> Self {
 24        Self {
 25            data: module.new_scalar_znx(rank),
 26            dist: SecretDistribution::NONE,
 27        }
 28    }
 29}
 30
 31impl<DataSelf> SecretKey<DataSelf> {
 32    pub fn n(&self) -> usize {
 33        self.data.n()
 34    }
 35
 36    pub fn log_n(&self) -> usize {
 37        self.data.log_n()
 38    }
 39
 40    pub fn rank(&self) -> usize {
 41        self.data.cols()
 42    }
 43}
 44
 45impl<S: AsMut<[u8]> + AsRef<[u8]>> SecretKey<S> {
 46    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 47        (0..self.rank()).for_each(|i| {
 48            self.data.fill_ternary_prob(i, prob, source);
 49        });
 50        self.dist = SecretDistribution::TernaryProb(prob);
 51    }
 52
 53    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 54        (0..self.rank()).for_each(|i| {
 55            self.data.fill_ternary_hw(i, hw, source);
 56        });
 57        self.dist = SecretDistribution::TernaryFixed(hw);
 58    }
 59
 60    pub fn fill_zero(&mut self) {
 61        self.data.zero();
 62        self.dist = SecretDistribution::ZERO;
 63    }
 64}
 65
 66pub struct SecretKeyFourier<T, B: Backend> {
 67    pub(crate) data: ScalarZnxDft<T, B>,
 68    pub(crate) dist: SecretDistribution,
 69}
 70
 71impl<DataSelf, B: Backend> SecretKeyFourier<DataSelf, B> {
 72    pub fn n(&self) -> usize {
 73        self.data.n()
 74    }
 75
 76    pub fn log_n(&self) -> usize {
 77        self.data.log_n()
 78    }
 79
 80    pub fn rank(&self) -> usize {
 81        self.data.cols()
 82    }
 83}
 84
 85impl<B: Backend> SecretKeyFourier<Vec<u8>, B> {
 86    pub fn alloc(module: &Module<B>, rank: usize) -> Self {
 87        Self {
 88            data: module.new_scalar_znx_dft(rank),
 89            dist: SecretDistribution::NONE,
 90        }
 91    }
 92}
 93
 94impl<D: AsRef<[u8]> + AsMut<[u8]>> SecretKeyFourier<D, FFT64> {
 95    pub fn dft<S: AsRef<[u8]>>(&mut self, module: &Module<FFT64>, sk: &SecretKey<S>) {
 96        #[cfg(debug_assertions)]
 97        {
 98            match sk.dist {
 99                SecretDistribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
100                _ => {}
101            }
102
103            assert_eq!(self.n(), module.n());
104            assert_eq!(sk.n(), module.n());
105            assert_eq!(self.rank(), sk.rank());
106        }
107
108        (0..self.rank()).for_each(|i| {
109            module.svp_prepare(&mut self.data, i, &sk.data, i);
110        });
111        self.dist = sk.dist;
112    }
113}
114
115pub struct GLWEPublicKey<D, B: Backend> {
116    pub(crate) data: GLWECiphertextFourier<D, B>,
117    pub(crate) dist: SecretDistribution,
118}
119
120impl<B: Backend> GLWEPublicKey<Vec<u8>, B> {
121    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
122        Self {
123            data: GLWECiphertextFourier::alloc(module, basek, k, rank),
124            dist: SecretDistribution::NONE,
125        }
126    }
127}
128
129impl<T, B: Backend> Infos for GLWEPublicKey<T, B> {
130    type Inner = VecZnxDft<T, B>;
131
132    fn inner(&self) -> &Self::Inner {
133        &self.data.data
134    }
135
136    fn basek(&self) -> usize {
137        self.data.basek
138    }
139
140    fn k(&self) -> usize {
141        self.data.k
142    }
143}
144
145impl<T, B: Backend> GLWEPublicKey<T, B> {
146    pub fn rank(&self) -> usize {
147        self.cols() - 1
148    }
149}
150
151impl<C: AsRef<[u8]> + AsMut<[u8]>> GLWEPublicKey<C, FFT64> {
152    pub fn generate_from_sk<S: AsRef<[u8]>>(
153        &mut self,
154        module: &Module<FFT64>,
155        sk_dft: &SecretKeyFourier<S, FFT64>,
156        source_xa: &mut Source,
157        source_xe: &mut Source,
158        sigma: f64,
159    ) {
160        #[cfg(debug_assertions)]
161        {
162            match sk_dft.dist {
163                SecretDistribution::NONE => panic!("invalid sk_dft: SecretDistribution::NONE"),
164                _ => {}
165            }
166        }
167
168        // Its ok to allocate scratch space here since pk is usually generated only once.
169        let mut scratch: ScratchOwned = ScratchOwned::new(GLWECiphertextFourier::encrypt_sk_scratch_space(
170            module,
171            self.rank(),
172            self.size(),
173        ));
174        self.data.encrypt_zero_sk(
175            module,
176            sk_dft,
177            source_xa,
178            source_xe,
179            sigma,
180            scratch.borrow(),
181        );
182        self.dist = sk_dft.dist;
183    }
184}