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, LWEInfos, LWEToGLWESwitchingKey, Rank, Rows, TorusPrecision,
  8    prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc},
  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> GGLWELayoutInfos for LWEToGLWESwitchingKeyPrepared<D, B> {
 40    fn digits(&self) -> Digits {
 41        self.0.digits()
 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 rows(&self) -> Rows {
 53        self.0.rows()
 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: GGLWELayoutInfos,
 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.digits().0,
 70            1,
 71            "digits > 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, rows: Rows, rank_out: Rank) -> Self
 77    where
 78        Module<B>: VmpPMatAlloc<B>,
 79    {
 80        Self(GGLWESwitchingKeyPrepared::alloc_with(
 81            module,
 82            base2k,
 83            k,
 84            rows,
 85            Digits(1),
 86            Rank(1),
 87            rank_out,
 88        ))
 89    }
 90
 91    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
 92    where
 93        A: GGLWELayoutInfos,
 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.digits().0,
103            1,
104            "digits > 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, rows: Rows, rank_out: Rank) -> usize
110    where
111        Module<B>: VmpPMatAllocBytes,
112    {
113        GGLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rows, Digits(1), Rank(1), rank_out)
114    }
115}
116
117impl<D: DataRef, B: Backend> PrepareAlloc<B, LWEToGLWESwitchingKeyPrepared<Vec<u8>, B>> for LWEToGLWESwitchingKey<D>
118where
119    Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
120{
121    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
122        let mut ksk_prepared: LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> = LWEToGLWESwitchingKeyPrepared::alloc(module, self);
123        ksk_prepared.prepare(module, self, scratch);
124        ksk_prepared
125    }
126}
127
128impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, LWEToGLWESwitchingKey<DR>> for LWEToGLWESwitchingKeyPrepared<DM, B>
129where
130    Module<B>: VmpPrepare<B>,
131{
132    fn prepare(&mut self, module: &Module<B>, other: &LWEToGLWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
133        self.0.prepare(module, &other.0, scratch);
134    }
135}