1use backend::hal::{
  2    api::{ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxDftAlloc, VecZnxDftAllocBytes, VecZnxDftFromVecZnx},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, Scratch, ScratchOwned, VecZnx, VecZnxDft, WriterTo},
  4    oep::{ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeVecZnxDftImpl, TakeVecZnxImpl},
  5};
  6use sampling::source::Source;
  7
  8use crate::{GLWECiphertext, GLWEEncryptSkFamily, GLWESecretExec, Infos, dist::Distribution};
  9
 10pub trait GLWEPublicKeyFamily<B: Backend> = GLWEEncryptSkFamily<B>;
 11
 12#[derive(PartialEq, Eq)]
 13pub struct GLWEPublicKey<D: Data> {
 14    pub(crate) data: VecZnx<D>,
 15    pub(crate) basek: usize,
 16    pub(crate) k: usize,
 17    pub(crate) dist: Distribution,
 18}
 19
 20impl GLWEPublicKey<Vec<u8>> {
 21    pub fn alloc(n: usize, basek: usize, k: usize, rank: usize) -> Self {
 22        Self {
 23            data: VecZnx::alloc(n, rank + 1, k.div_ceil(basek)),
 24            basek: basek,
 25            k: k,
 26            dist: Distribution::NONE,
 27        }
 28    }
 29
 30    pub fn bytes_of(n: usize, basek: usize, k: usize, rank: usize) -> usize {
 31        VecZnx::alloc_bytes(n, rank + 1, k.div_ceil(basek))
 32    }
 33}
 34
 35impl<D: Data> Infos for GLWEPublicKey<D> {
 36    type Inner = VecZnx<D>;
 37
 38    fn inner(&self) -> &Self::Inner {
 39        &self.data
 40    }
 41
 42    fn basek(&self) -> usize {
 43        self.basek
 44    }
 45
 46    fn k(&self) -> usize {
 47        self.k
 48    }
 49}
 50
 51impl<D: Data> GLWEPublicKey<D> {
 52    pub fn rank(&self) -> usize {
 53        self.cols() - 1
 54    }
 55}
 56
 57impl<D: DataMut> GLWEPublicKey<D> {
 58    pub fn generate_from_sk<S: DataRef, B: Backend>(
 59        &mut self,
 60        module: &Module<B>,
 61        sk: &GLWESecretExec<S, B>,
 62        source_xa: &mut Source,
 63        source_xe: &mut Source,
 64        sigma: f64,
 65    ) where
 66        Module<B>: GLWEPublicKeyFamily<B>,
 67        B: ScratchOwnedAllocImpl<B>
 68            + ScratchOwnedBorrowImpl<B>
 69            + TakeVecZnxDftImpl<B>
 70            + ScratchAvailableImpl<B>
 71            + TakeVecZnxImpl<B>,
 72    {
 73        #[cfg(debug_assertions)]
 74        {
 75            assert_eq!(self.n(), sk.n());
 76
 77            match sk.dist {
 78                Distribution::NONE => panic!("invalid sk: SecretDistribution::NONE"),
 79                _ => {}
 80            }
 81        }
 82
 83        // Its ok to allocate scratch space here since pk is usually generated only once.
 84        let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWECiphertext::encrypt_sk_scratch_space(
 85            module,
 86            self.n(),
 87            self.basek(),
 88            self.k(),
 89        ));
 90
 91        let mut tmp: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(self.n(), self.basek(), self.k(), self.rank());
 92        tmp.encrypt_zero_sk(module, sk, source_xa, source_xe, sigma, scratch.borrow());
 93        self.dist = sk.dist;
 94    }
 95}
 96
 97use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 98
 99impl<D: DataMut> ReaderFrom for GLWEPublicKey<D> {
100    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
101        self.k = reader.read_u64::<LittleEndian>()? as usize;
102        self.basek = reader.read_u64::<LittleEndian>()? as usize;
103        match Distribution::read_from(reader) {
104            Ok(dist) => self.dist = dist,
105            Err(e) => return Err(e),
106        }
107        self.data.read_from(reader)
108    }
109}
110
111impl<D: DataRef> WriterTo for GLWEPublicKey<D> {
112    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
113        writer.write_u64::<LittleEndian>(self.k as u64)?;
114        writer.write_u64::<LittleEndian>(self.basek as u64)?;
115        match self.dist.write_to(writer) {
116            Ok(()) => {}
117            Err(e) => return Err(e),
118        }
119        self.data.write_to(writer)
120    }
121}
122
123#[derive(PartialEq, Eq)]
124pub struct GLWEPublicKeyExec<D: Data, B: Backend> {
125    pub(crate) data: VecZnxDft<D, B>,
126    pub(crate) basek: usize,
127    pub(crate) k: usize,
128    pub(crate) dist: Distribution,
129}
130
131impl<D: Data, B: Backend> Infos for GLWEPublicKeyExec<D, B> {
132    type Inner = VecZnxDft<D, B>;
133
134    fn inner(&self) -> &Self::Inner {
135        &self.data
136    }
137
138    fn basek(&self) -> usize {
139        self.basek
140    }
141
142    fn k(&self) -> usize {
143        self.k
144    }
145}
146
147impl<D: Data, B: Backend> GLWEPublicKeyExec<D, B> {
148    pub fn rank(&self) -> usize {
149        self.cols() - 1
150    }
151}
152
153impl<B: Backend> GLWEPublicKeyExec<Vec<u8>, B> {
154    pub fn alloc(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> Self
155    where
156        Module<B>: VecZnxDftAlloc<B>,
157    {
158        Self {
159            data: module.vec_znx_dft_alloc(n, rank + 1, k.div_ceil(basek)),
160            basek: basek,
161            k: k,
162            dist: Distribution::NONE,
163        }
164    }
165
166    pub fn bytes_of(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
167    where
168        Module<B>: VecZnxDftAllocBytes,
169    {
170        module.vec_znx_dft_alloc_bytes(n, rank + 1, k.div_ceil(basek))
171    }
172
173    pub fn from<DataOther>(module: &Module<B>, other: &GLWEPublicKey<DataOther>, scratch: &mut Scratch<B>) -> Self
174    where
175        DataOther: DataRef,
176        Module<B>: VecZnxDftAlloc<B> + VecZnxDftFromVecZnx<B>,
177    {
178        let mut pk_exec: GLWEPublicKeyExec<Vec<u8>, B> =
179            GLWEPublicKeyExec::alloc(module, other.n(), other.basek(), other.k(), other.rank());
180        pk_exec.prepare(module, other, scratch);
181        pk_exec
182    }
183}
184
185impl<D: DataMut, B: Backend> GLWEPublicKeyExec<D, B> {
186    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWEPublicKey<DataOther>, _scratch: &mut Scratch<B>)
187    where
188        DataOther: DataRef,
189        Module<B>: VecZnxDftFromVecZnx<B>,
190    {
191        #[cfg(debug_assertions)]
192        {
193            assert_eq!(self.n(), other.n());
194            assert_eq!(self.size(), other.size());
195        }
196
197        (0..self.cols()).for_each(|i| {
198            module.vec_znx_dft_from_vec_znx(1, 0, &mut self.data, i, &other.data, i);
199        });
200        self.k = other.k;
201        self.basek = other.basek;
202        self.dist = other.dist;
203    }
204}