1use base2k::{
 2    Backend, Module, Scalar, ScalarAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScalarZnxDftToMut, Scratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftToMut, ZnxInfos, FFT64
 3};
 4use sampling::source::Source;
 5
 6use crate::elem::derive_size;
 7
 8pub struct SecretKey<T> {
 9    data: T,
10}
11
12impl<T> SecretKey<T> {
13    pub fn data(&self) -> &T {
14        &self.data
15    }
16
17    pub fn data_mut(&mut self) -> &mut T {
18        &mut self.data
19    }
20}
21
22impl SecretKey<Scalar<Vec<u8>>> {
23    pub fn new<B: Backend>(module: &Module<B>) -> Self {
24        Self {
25            data: module.new_scalar(1),
26        }
27    }
28
29    pub fn fill_ternary_prob(&mut self, prob: f64, source: &mut Source) {
30        self.data.fill_ternary_prob(0, prob, source);
31    }
32
33    pub fn fill_ternary_hw(&mut self, hw: usize, source: &mut Source) {
34        self.data.fill_ternary_hw(0, hw, source);
35    }
36
37    pub fn svp_prepare<D>(&self, module: &Module<FFT64>, sk_prep: &mut SecretKey<ScalarZnxDft<D, FFT64>>)
38    where
39        ScalarZnxDft<D, base2k::FFT64>: ScalarZnxDftToMut<base2k::FFT64>,
40    {
41        module.svp_prepare(&mut sk_prep.data, 0, &self.data, 0)
42    }
43}
44
45type SecretKeyPrep<C, B> = SecretKey<ScalarZnxDft<C, B>>;
46
47impl<B: Backend> SecretKey<ScalarZnxDft<Vec<u8>, B>> {
48    pub fn new(module: &Module<B>) -> Self{
49        Self{
50            data: module.new_scalar_znx_dft(1)
51        }
52    }
53}
54
55pub struct PublicKey<D, B: Backend> {
56    data: VecZnxDft<D, B>,
57}
58
59impl<B: Backend> PublicKey<Vec<u8>, B> {
60    pub fn new(module: &Module<B>, log_base2k: usize, log_q: usize) -> Self {
61        Self {
62            data: module.new_vec_znx_dft(2, derive_size(log_base2k, log_q)),
63        }
64    }
65}
66
67impl<B: Backend, D: VecZnxDftToMut<B>> PublicKey<D, B> {
68    pub fn generate<S>(&mut self, module: &Module<B>, sk: &SecretKey<ScalarZnxDft<S, B>>, scratch: &mut Scratch)
69    where
70        ScalarZnxDft<S, B>: ScalarZnxDftToMut<B>,
71    {
72    }
73}