1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWESwitchingKey, Rank, TorusPrecision,
  8    prepared::GLWESwitchingKeyPrepared,
  9};
 10
 11#[derive(PartialEq, Eq)]
 12pub struct LWESwitchingKeyPrepared<D: Data, B: Backend>(pub(crate) GLWESwitchingKeyPrepared<D, B>);
 13
 14impl<D: Data, B: Backend> LWEInfos for LWESwitchingKeyPrepared<D, B> {
 15    fn base2k(&self) -> Base2K {
 16        self.0.base2k()
 17    }
 18
 19    fn k(&self) -> TorusPrecision {
 20        self.0.k()
 21    }
 22
 23    fn n(&self) -> Degree {
 24        self.0.n()
 25    }
 26
 27    fn size(&self) -> usize {
 28        self.0.size()
 29    }
 30}
 31impl<D: Data, B: Backend> GLWEInfos for LWESwitchingKeyPrepared<D, B> {
 32    fn rank(&self) -> Rank {
 33        self.rank_out()
 34    }
 35}
 36
 37impl<D: Data, B: Backend> GGLWEInfos for LWESwitchingKeyPrepared<D, B> {
 38    fn dsize(&self) -> Dsize {
 39        self.0.dsize()
 40    }
 41
 42    fn rank_in(&self) -> Rank {
 43        self.0.rank_in()
 44    }
 45
 46    fn rank_out(&self) -> Rank {
 47        self.0.rank_out()
 48    }
 49
 50    fn dnum(&self) -> Dnum {
 51        self.0.dnum()
 52    }
 53}
 54
 55impl<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
 56    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 57    where
 58        A: GGLWEInfos,
 59        Module<B>: VmpPMatAlloc<B>,
 60    {
 61        debug_assert_eq!(
 62            infos.dsize().0,
 63            1,
 64            "dsize > 1 is not supported for LWESwitchingKey"
 65        );
 66        debug_assert_eq!(
 67            infos.rank_in().0,
 68            1,
 69            "rank_in > 1 is not supported for LWESwitchingKey"
 70        );
 71        debug_assert_eq!(
 72            infos.rank_out().0,
 73            1,
 74            "rank_out > 1 is not supported for LWESwitchingKey"
 75        );
 76        Self(GLWESwitchingKeyPrepared::alloc(module, infos))
 77    }
 78
 79    pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum) -> Self
 80    where
 81        Module<B>: VmpPMatAlloc<B>,
 82    {
 83        Self(GLWESwitchingKeyPrepared::alloc_with(
 84            module,
 85            base2k,
 86            k,
 87            Rank(1),
 88            Rank(1),
 89            dnum,
 90            Dsize(1),
 91        ))
 92    }
 93
 94    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
 95    where
 96        A: GGLWEInfos,
 97        Module<B>: VmpPMatAllocBytes,
 98    {
 99        debug_assert_eq!(
100            infos.dsize().0,
101            1,
102            "dsize > 1 is not supported for LWESwitchingKey"
103        );
104        debug_assert_eq!(
105            infos.rank_in().0,
106            1,
107            "rank_in > 1 is not supported for LWESwitchingKey"
108        );
109        debug_assert_eq!(
110            infos.rank_out().0,
111            1,
112            "rank_out > 1 is not supported for LWESwitchingKey"
113        );
114        GLWESwitchingKeyPrepared::alloc_bytes(module, infos)
115    }
116
117    pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum) -> usize
118    where
119        Module<B>: VmpPMatAllocBytes,
120    {
121        GLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, Rank(1), Rank(1), dnum, Dsize(1))
122    }
123}
124
125impl<B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for LWESwitchingKeyPrepared<Vec<u8>, B>
126where
127    GLWESwitchingKeyPrepared<Vec<u8>, B>: PrepareScratchSpace<B, A>,
128{
129    fn prepare_scratch_space(module: &Module<B>, infos: &A) -> usize {
130        GLWESwitchingKeyPrepared::prepare_scratch_space(module, infos)
131    }
132}
133
134impl<D: DataRef, B: Backend> PrepareAlloc<B, LWESwitchingKeyPrepared<Vec<u8>, B>> for LWESwitchingKey<D>
135where
136    Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
137{
138    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> LWESwitchingKeyPrepared<Vec<u8>, B> {
139        let mut ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, B> = LWESwitchingKeyPrepared::alloc(module, self);
140        ksk_prepared.prepare(module, self, scratch);
141        ksk_prepared
142    }
143}
144
145impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, LWESwitchingKey<DR>> for LWESwitchingKeyPrepared<DM, B>
146where
147    Module<B>: VmpPrepare<B>,
148{
149    fn prepare(&mut self, module: &Module<B>, other: &LWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
150        self.0.prepare(module, &other.0, scratch);
151    }
152}