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