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