1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Digits, GGLWELayoutInfos, GLWEInfos, GLWEToLWESwitchingKey, LWEInfos, Rank, Rows, TorusPrecision,
  8    prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc},
  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> GGLWELayoutInfos for GLWEToLWESwitchingKeyPrepared<D, B> {
 39    fn rank_in(&self) -> Rank {
 40        self.0.rank_in()
 41    }
 42
 43    fn digits(&self) -> Digits {
 44        self.0.digits()
 45    }
 46
 47    fn rank_out(&self) -> Rank {
 48        self.0.rank_out()
 49    }
 50
 51    fn rows(&self) -> Rows {
 52        self.0.rows()
 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: GGLWELayoutInfos,
 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.digits().0,
 69            1,
 70            "digits > 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, rows: Rows, rank_in: Rank) -> Self
 76    where
 77        Module<B>: VmpPMatAlloc<B>,
 78    {
 79        Self(GGLWESwitchingKeyPrepared::alloc_with(
 80            module,
 81            base2k,
 82            k,
 83            rows,
 84            Digits(1),
 85            rank_in,
 86            Rank(1),
 87        ))
 88    }
 89
 90    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
 91    where
 92        A: GGLWELayoutInfos,
 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.digits().0,
102            1,
103            "digits > 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, rows: Rows, rank_in: Rank) -> usize
109    where
110        Module<B>: VmpPMatAllocBytes,
111    {
112        GGLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rows, Digits(1), rank_in, Rank(1))
113    }
114}
115
116impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWEToLWESwitchingKeyPrepared<Vec<u8>, B>> for GLWEToLWESwitchingKey<D>
117where
118    Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
119{
120    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
121        let mut ksk_prepared: GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> = GLWEToLWESwitchingKeyPrepared::alloc(module, self);
122        ksk_prepared.prepare(module, self, scratch);
123        ksk_prepared
124    }
125}
126
127impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWEToLWESwitchingKey<DR>> for GLWEToLWESwitchingKeyPrepared<DM, B>
128where
129    Module<B>: VmpPrepare<B>,
130{
131    fn prepare(&mut self, module: &Module<B>, other: &GLWEToLWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
132        self.0.prepare(module, &other.0, scratch);
133    }
134}