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