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