1use poulpy_hal::{
  2    api::{SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, SvpPPol, SvpPPolToMut, SvpPPolToRef, ZnxInfos},
  4};
  5
  6use crate::{
  7    dist::Distribution,
  8    layouts::{
  9        Base2K, Degree, GLWEInfos, GLWESecret, LWEInfos, Rank, TorusPrecision,
 10        prepared::{Prepare, PrepareAlloc, PrepareScratchSpace},
 11    },
 12};
 13
 14pub struct GLWESecretPrepared<D: Data, B: Backend> {
 15    pub(crate) data: SvpPPol<D, B>,
 16    pub(crate) dist: Distribution,
 17}
 18
 19impl<D: Data, B: Backend> LWEInfos for GLWESecretPrepared<D, B> {
 20    fn base2k(&self) -> Base2K {
 21        Base2K(0)
 22    }
 23
 24    fn k(&self) -> TorusPrecision {
 25        TorusPrecision(0)
 26    }
 27
 28    fn n(&self) -> Degree {
 29        Degree(self.data.n() as u32)
 30    }
 31
 32    fn size(&self) -> usize {
 33        self.data.size()
 34    }
 35}
 36impl<D: Data, B: Backend> GLWEInfos for GLWESecretPrepared<D, B> {
 37    fn rank(&self) -> Rank {
 38        Rank(self.data.cols() as u32)
 39    }
 40}
 41impl<B: Backend> GLWESecretPrepared<Vec<u8>, B> {
 42    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 43    where
 44        A: GLWEInfos,
 45        Module<B>: SvpPPolAlloc<B>,
 46    {
 47        assert_eq!(module.n() as u32, infos.n());
 48        Self::alloc_with(module, infos.rank())
 49    }
 50
 51    pub fn alloc_with(module: &Module<B>, rank: Rank) -> Self
 52    where
 53        Module<B>: SvpPPolAlloc<B>,
 54    {
 55        Self {
 56            data: module.svp_ppol_alloc(rank.into()),
 57            dist: Distribution::NONE,
 58        }
 59    }
 60
 61    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
 62    where
 63        A: GLWEInfos,
 64        Module<B>: SvpPPolAllocBytes,
 65    {
 66        assert_eq!(module.n() as u32, infos.n());
 67        Self::alloc_bytes_with(module, infos.rank())
 68    }
 69
 70    pub fn alloc_bytes_with(module: &Module<B>, rank: Rank) -> usize
 71    where
 72        Module<B>: SvpPPolAllocBytes,
 73    {
 74        module.svp_ppol_alloc_bytes(rank.into())
 75    }
 76}
 77
 78impl<D: Data, B: Backend> GLWESecretPrepared<D, B> {
 79    pub fn n(&self) -> Degree {
 80        Degree(self.data.n() as u32)
 81    }
 82
 83    pub fn rank(&self) -> Rank {
 84        Rank(self.data.cols() as u32)
 85    }
 86}
 87
 88impl<B: Backend, A: GLWEInfos> PrepareScratchSpace<B, A> for GLWESecretPrepared<Vec<u8>, B> {
 89    fn prepare_scratch_space(_module: &Module<B>, _infos: &A) -> usize {
 90        0
 91    }
 92}
 93
 94impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWESecretPrepared<Vec<u8>, B>> for GLWESecret<D>
 95where
 96    Module<B>: SvpPrepare<B> + SvpPPolAlloc<B>,
 97{
 98    fn prepare_alloc(&self, module: &Module<B>, _scratch: &mut Scratch<B>) -> GLWESecretPrepared<Vec<u8>, B> {
 99        let mut sk_dft: GLWESecretPrepared<Vec<u8>, B> = GLWESecretPrepared::alloc(module, self);
100        sk_dft.prepare(module, self, _scratch);
101        sk_dft
102    }
103}
104
105impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWESecret<DR>> for GLWESecretPrepared<DM, B>
106where
107    Module<B>: SvpPrepare<B>,
108{
109    fn prepare(&mut self, module: &Module<B>, other: &GLWESecret<DR>, _scratch: &mut Scratch<B>) {
110        (0..self.rank().into()).for_each(|i| {
111            module.svp_prepare(&mut self.data, i, &other.data, i);
112        });
113        self.dist = other.dist
114    }
115}
116
117pub trait GLWESecretPreparedToRef<B: Backend> {
118    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B>;
119}
120
121impl<D: DataRef, B: Backend> GLWESecretPreparedToRef<B> for GLWESecretPrepared<D, B> {
122    fn to_ref(&self) -> GLWESecretPrepared<&[u8], B> {
123        GLWESecretPrepared {
124            data: self.data.to_ref(),
125            dist: self.dist,
126        }
127    }
128}
129
130pub trait GLWESecretPreparedToMut<B: Backend> {
131    fn to_ref(&mut self) -> GLWESecretPrepared<&mut [u8], B>;
132}
133
134impl<D: DataMut, B: Backend> GLWESecretPreparedToMut<B> for GLWESecretPrepared<D, B> {
135    fn to_ref(&mut self) -> GLWESecretPrepared<&mut [u8], B> {
136        GLWESecretPrepared {
137            dist: self.dist,
138            data: self.data.to_mut(),
139        }
140    }
141}