1use backend::{
  2    Backend, FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScratchOwned, VecZnxDft,
  3    ZnxInfos, ZnxZero,
  4};
  5use sampling::source::Source;
  6
  7use crate::{GLWECiphertextFourier, Infos};
  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 GLWESecret<T, B: Backend> {
 18    pub(crate) data: ScalarZnx<T>,
 19    pub(crate) data_fourier: ScalarZnxDft<T, B>,
 20    pub(crate) dist: SecretDistribution,
 21}
 22
 23impl<B: Backend> GLWESecret<Vec<u8>, B> {
 24    pub fn alloc(module: &Module<B>, rank: usize) -> Self {
 25        Self {
 26            data: module.new_scalar_znx(rank),
 27            data_fourier: module.new_scalar_znx_dft(rank),
 28            dist: SecretDistribution::NONE,
 29        }
 30    }
 31
 32    pub fn bytes_of(module: &Module<B>, rank: usize) -> usize {
 33        module.bytes_of_scalar_znx(rank) + module.bytes_of_scalar_znx_dft(rank)
 34    }
 35}
 36
 37impl<DataSelf, B: Backend> GLWESecret<DataSelf, B> {
 38    pub fn n(&self) -> usize {
 39        self.data.n()
 40    }
 41
 42    pub fn log_n(&self) -> usize {
 43        self.data.log_n()
 44    }
 45
 46    pub fn rank(&self) -> usize {
 47        self.data.cols()
 48    }
 49}
 50
 51impl<S: AsMut<[u8]> + AsRef<[u8]>> GLWESecret<S, FFT64> {
 52    pub fn fill_ternary_prob(&mut self, module: &Module<FFT64>, prob: f64, source: &mut Source) {
 53        (0..self.rank()).for_each(|i| {
 54            self.data.fill_ternary_prob(i, prob, source);
 55        });
 56        self.prep_fourier(module);
 57        self.dist = SecretDistribution::TernaryProb(prob);
 58    }
 59
 60    pub fn fill_ternary_hw(&mut self, module: &Module<FFT64>, hw: usize, source: &mut Source) {
 61        (0..self.rank()).for_each(|i| {
 62            self.data.fill_ternary_hw(i, hw, source);
 63        });
 64        self.prep_fourier(module);
 65        self.dist = SecretDistribution::TernaryFixed(hw);
 66    }
 67
 68    pub fn fill_zero(&mut self) {
 69        self.data.zero();
 70        self.dist = SecretDistribution::ZERO;
 71    }
 72
 73    pub(crate) fn prep_fourier(&mut self, module: &Module<FFT64>) {
 74        (0..self.rank()).for_each(|i| {
 75            module.svp_prepare(&mut self.data_fourier, i, &self.data, i);
 76        });
 77    }
 78}
 79
 80pub struct GLWEPublicKey<D, B: Backend> {
 81    pub(crate) data: GLWECiphertextFourier<D, B>,
 82    pub(crate) dist: SecretDistribution,
 83}
 84
 85impl<B: Backend> GLWEPublicKey<Vec<u8>, B> {
 86    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
 87        Self {
 88            data: GLWECiphertextFourier::alloc(module, basek, k, rank),
 89            dist: SecretDistribution::NONE,
 90        }
 91    }
 92
 93    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize {
 94        GLWECiphertextFourier::<Vec<u8>, B>::bytes_of(module, basek, k, rank)
 95    }
 96}
 97
 98impl<T, B: Backend> Infos for GLWEPublicKey<T, B> {
 99    type Inner = VecZnxDft<T, B>;
100
101    fn inner(&self) -> &Self::Inner {
102        &self.data.data
103    }
104
105    fn basek(&self) -> usize {
106        self.data.basek
107    }
108
109    fn k(&self) -> usize {
110        self.data.k
111    }
112}
113
114impl<T, B: Backend> GLWEPublicKey<T, B> {
115    pub fn rank(&self) -> usize {
116        self.cols() - 1
117    }
118}
119
120impl<C: AsRef<[u8]> + AsMut<[u8]>> GLWEPublicKey<C, FFT64> {
121    pub fn generate_from_sk<S: AsRef<[u8]>>(
122        &mut self,
123        module: &Module<FFT64>,
124        sk: &GLWESecret<S, FFT64>,
125        source_xa: &mut Source,
126        source_xe: &mut Source,
127        sigma: f64,
128    ) {
129        #[cfg(debug_assertions)]
130        {
131            match sk.dist {
132                SecretDistribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
133                _ => {}
134            }
135        }
136
137        // Its ok to allocate scratch space here since pk is usually generated only once.
138        let mut scratch: ScratchOwned = ScratchOwned::new(GLWECiphertextFourier::encrypt_sk_scratch_space(
139            module,
140            self.basek(),
141            self.k(),
142            self.rank(),
143        ));
144
145        self.data
146            .encrypt_zero_sk(module, sk, source_xa, source_xe, sigma, scratch.borrow());
147        self.dist = sk.dist;
148    }
149}