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