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