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