1use base2k::{
  2    Backend, FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScalarZnxDftToMut,
  3    ScalarZnxDftToRef, ScalarZnxToMut, ScalarZnxToRef, Scratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftToMut,
  4};
  5use sampling::source::Source;
  6
  7use crate::elem::derive_size;
  8
  9pub struct SecretKey<T> {
 10    pub data: ScalarZnx<T>,
 11}
 12
 13impl SecretKey<Vec<u8>> {
 14    pub fn new<B: Backend>(module: &Module<B>) -> Self {
 15        Self {
 16            data: module.new_scalar(1),
 17        }
 18    }
 19}
 20
 21impl<S> SecretKey<S>
 22where
 23    S: AsMut<[u8]> + AsRef<[u8]>,
 24{
 25    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
 26        self.data.fill_ternary_prob(0, prob, source);
 27    }
 28
 29    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
 30        self.data.fill_ternary_hw(0, hw, source);
 31    }
 32}
 33
 34impl<C> ScalarZnxToMut for SecretKey<C>
 35where
 36    ScalarZnx<C>: ScalarZnxToMut,
 37{
 38    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
 39        self.data.to_mut()
 40    }
 41}
 42
 43impl<C> ScalarZnxToRef for SecretKey<C>
 44where
 45    ScalarZnx<C>: ScalarZnxToRef,
 46{
 47    fn to_ref(&self) -> ScalarZnx<&[u8]> {
 48        self.data.to_ref()
 49    }
 50}
 51
 52pub struct SecretKeyDft<T, B: Backend> {
 53    pub data: ScalarZnxDft<T, B>,
 54}
 55
 56impl<B: Backend> SecretKeyDft<Vec<u8>, B> {
 57    pub fn new(module: &Module<B>) -> Self {
 58        Self {
 59            data: module.new_scalar_znx_dft(1),
 60        }
 61    }
 62
 63    pub fn dft<S>(&mut self, module: &Module<FFT64>, sk: &SecretKey<S>)
 64    where
 65        SecretKeyDft<Vec<u8>, B>: ScalarZnxDftToMut<base2k::FFT64>,
 66        SecretKey<S>: ScalarZnxToRef,
 67    {
 68        module.svp_prepare(self, 0, sk, 0)
 69    }
 70}
 71
 72impl<C, B: Backend> ScalarZnxDftToMut<B> for SecretKeyDft<C, B>
 73where
 74    ScalarZnxDft<C, B>: ScalarZnxDftToMut<B>,
 75{
 76    fn to_mut(&mut self) -> ScalarZnxDft<&mut [u8], B> {
 77        self.data.to_mut()
 78    }
 79}
 80
 81impl<C, B: Backend> ScalarZnxDftToRef<B> for SecretKeyDft<C, B>
 82where
 83    ScalarZnxDft<C, B>: ScalarZnxDftToRef<B>,
 84{
 85    fn to_ref(&self) -> ScalarZnxDft<&[u8], B> {
 86        self.data.to_ref()
 87    }
 88}
 89
 90pub struct PublicKey<D, B: Backend> {
 91    pub data: VecZnxDft<D, B>,
 92}
 93
 94impl<B: Backend> PublicKey<Vec<u8>, B> {
 95    pub fn new(module: &Module<B>, log_base2k: usize, log_q: usize) -> Self {
 96        Self {
 97            data: module.new_vec_znx_dft(2, derive_size(log_base2k, log_q)),
 98        }
 99    }
100}
101
102impl<B: Backend, D: VecZnxDftToMut<B>> PublicKey<D, B> {
103    pub fn generate<S>(&mut self, module: &Module<B>, sk: &SecretKey<ScalarZnxDft<S, B>>, scratch: &mut Scratch)
104    where
105        ScalarZnxDft<S, B>: ScalarZnxDftToMut<B>,
106    {
107    }
108}