1use base2k::{
  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, rlwe::RLWECtDft};
  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 new<B: Backend>(module: &Module<B>) -> Self {
 25        Self {
 26            data: module.new_scalar_znx(1),
 27            dist: SecretDistribution::NONE,
 28        }
 29    }
 30}
 31
 32impl<S> SecretKey<S>
 33where
 34    S: AsMut<[u8]> + AsRef<[u8]>,
 35{
 36    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 37        self.data.fill_ternary_prob(0, prob, source);
 38        self.dist = SecretDistribution::TernaryProb(prob);
 39    }
 40
 41    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 42        self.data.fill_ternary_hw(0, hw, source);
 43        self.dist = SecretDistribution::TernaryFixed(hw);
 44    }
 45
 46    pub fn fill_zero(&mut self) {
 47        self.data.zero();
 48        self.dist = SecretDistribution::ZERO;
 49    }
 50}
 51
 52impl<C> ScalarZnxToMut for SecretKey<C>
 53where
 54    ScalarZnx<C>: ScalarZnxToMut,
 55{
 56    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
 57        self.data.to_mut()
 58    }
 59}
 60
 61impl<C> ScalarZnxToRef for SecretKey<C>
 62where
 63    ScalarZnx<C>: ScalarZnxToRef,
 64{
 65    fn to_ref(&self) -> ScalarZnx<&[u8]> {
 66        self.data.to_ref()
 67    }
 68}
 69
 70pub struct SecretKeyDft<T, B: Backend> {
 71    pub data: ScalarZnxDft<T, B>,
 72    pub dist: SecretDistribution,
 73}
 74
 75impl<B: Backend> SecretKeyDft<Vec<u8>, B> {
 76    pub fn new(module: &Module<B>) -> Self {
 77        Self {
 78            data: module.new_scalar_znx_dft(1),
 79            dist: SecretDistribution::NONE,
 80        }
 81    }
 82
 83    pub fn dft<S>(&mut self, module: &Module<FFT64>, sk: &SecretKey<S>)
 84    where
 85        SecretKeyDft<Vec<u8>, B>: ScalarZnxDftToMut<base2k::FFT64>,
 86        SecretKey<S>: ScalarZnxToRef,
 87    {
 88        #[cfg(debug_assertions)]
 89        {
 90            match sk.dist {
 91                SecretDistribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
 92                _ => {}
 93            }
 94        }
 95
 96        module.svp_prepare(self, 0, sk, 0);
 97        self.dist = sk.dist;
 98    }
 99}
100
101impl<C, B: Backend> ScalarZnxDftToMut<B> for SecretKeyDft<C, B>
102where
103    ScalarZnxDft<C, B>: ScalarZnxDftToMut<B>,
104{
105    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
106        self.data.to_mut()
107    }
108}
109
110impl<C, B: Backend> ScalarZnxDftToRef<B> for SecretKeyDft<C, B>
111where
112    ScalarZnxDft<C, B>: ScalarZnxDftToRef<B>,
113{
114    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
115        self.data.to_ref()
116    }
117}
118
119pub struct PublicKey<D, B: Backend> {
120    pub data: RLWECtDft<D, B>,
121    pub dist: SecretDistribution,
122}
123
124impl<B: Backend> PublicKey<Vec<u8>, B> {
125    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
126        Self {
127            data: RLWECtDft::new(module, log_base2k, log_k),
128            dist: SecretDistribution::NONE,
129        }
130    }
131}
132
133impl<T, B: Backend> Infos for PublicKey<T, B> {
134    type Inner = VecZnxDft<T, B>;
135
136    fn inner(&self) -> &Self::Inner {
137        &self.data.data
138    }
139
140    fn log_base2k(&self) -> usize {
141        self.data.log_base2k
142    }
143
144    fn log_k(&self) -> usize {
145        self.data.log_k
146    }
147}
148
149impl<C, B: Backend> VecZnxDftToMut<B> for PublicKey<C, B>
150where
151    VecZnxDft<C, B>: VecZnxDftToMut<B>,
152{
153    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
154        self.data.to_mut()
155    }
156}
157
158impl<C, B: Backend> VecZnxDftToRef<B> for PublicKey<C, B>
159where
160    VecZnxDft<C, B>: VecZnxDftToRef<B>,
161{
162    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
163        self.data.to_ref()
164    }
165}
166
167impl<C> PublicKey<C, FFT64> {
168    pub fn generate<S>(
169        &mut self,
170        module: &Module<FFT64>,
171        sk_dft: &SecretKeyDft<S, FFT64>,
172        source_xa: &mut Source,
173        source_xe: &mut Source,
174        sigma: f64,
175        bound: f64,
176    ) where
177        VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
178        ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64> + ZnxInfos,
179    {
180        #[cfg(debug_assertions)]
181        {
182            match sk_dft.dist {
183                SecretDistribution::NONE => panic!("invalid sk_dft: SecretDistribution::NONE"),
184                _ => {}
185            }
186        }
187
188        // Its ok to allocate scratch space here since pk is usually generated only once.
189        let mut scratch: ScratchOwned = ScratchOwned::new(RLWECtDft::encrypt_zero_sk_scratch_space(
190            module,
191            self.size(),
192        ));
193        self.data.encrypt_zero_sk(
194            module,
195            sk_dft,
196            source_xa,
197            source_xe,
198            sigma,
199            bound,
200            scratch.borrow(),
201        );
202        self.dist = sk_dft.dist;
203    }
204}