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