1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat},
  4};
  5
  6use crate::layouts::{
  7    GGLWESwitchingKey, Infos,
  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<B: Backend> GGLWESwitchingKeyPrepared<Vec<u8>, B> {
 19    #[allow(clippy::too_many_arguments)]
 20    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self
 21    where
 22        Module<B>: VmpPMatAlloc<B>,
 23    {
 24        GGLWESwitchingKeyPrepared::<Vec<u8>, B> {
 25            key: GGLWECiphertextPrepared::alloc(module, basek, k, rows, digits, rank_in, rank_out),
 26            sk_in_n: 0,
 27            sk_out_n: 0,
 28        }
 29    }
 30
 31    #[allow(clippy::too_many_arguments)]
 32    pub fn bytes_of(
 33        module: &Module<B>,
 34        basek: usize,
 35        k: usize,
 36        rows: usize,
 37        digits: usize,
 38        rank_in: usize,
 39        rank_out: usize,
 40    ) -> usize
 41    where
 42        Module<B>: VmpPMatAllocBytes,
 43    {
 44        GGLWECiphertextPrepared::bytes_of(module, basek, k, rows, digits, rank_in, rank_out)
 45    }
 46}
 47
 48impl<D: Data, B: Backend> Infos for GGLWESwitchingKeyPrepared<D, B> {
 49    type Inner = VmpPMat<D, B>;
 50
 51    fn inner(&self) -> &Self::Inner {
 52        self.key.inner()
 53    }
 54
 55    fn basek(&self) -> usize {
 56        self.key.basek()
 57    }
 58
 59    fn k(&self) -> usize {
 60        self.key.k()
 61    }
 62}
 63
 64impl<D: Data, B: Backend> GGLWESwitchingKeyPrepared<D, B> {
 65    pub fn rank(&self) -> usize {
 66        self.key.data.cols_out() - 1
 67    }
 68
 69    pub fn rank_in(&self) -> usize {
 70        self.key.data.cols_in()
 71    }
 72
 73    pub fn rank_out(&self) -> usize {
 74        self.key.data.cols_out() - 1
 75    }
 76
 77    pub fn digits(&self) -> usize {
 78        self.key.digits()
 79    }
 80
 81    pub fn sk_degree_in(&self) -> usize {
 82        self.sk_in_n
 83    }
 84
 85    pub fn sk_degree_out(&self) -> usize {
 86        self.sk_out_n
 87    }
 88}
 89
 90impl<D: DataMut, DR: DataRef, B: Backend> Prepare<B, GGLWESwitchingKey<DR>> for GGLWESwitchingKeyPrepared<D, B>
 91where
 92    Module<B>: VmpPrepare<B>,
 93{
 94    fn prepare(&mut self, module: &Module<B>, other: &GGLWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
 95        self.key.prepare(module, &other.key, scratch);
 96        self.sk_in_n = other.sk_in_n;
 97        self.sk_out_n = other.sk_out_n;
 98    }
 99}
100
101impl<D: DataRef, B: Backend> PrepareAlloc<B, GGLWESwitchingKeyPrepared<Vec<u8>, B>> for GGLWESwitchingKey<D>
102where
103    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
104{
105    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GGLWESwitchingKeyPrepared<Vec<u8>, B> {
106        let mut atk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = GGLWESwitchingKeyPrepared::alloc(
107            module,
108            self.basek(),
109            self.k(),
110            self.rows(),
111            self.digits(),
112            self.rank_in(),
113            self.rank_out(),
114        );
115        atk_prepared.prepare(module, self, scratch);
116        atk_prepared
117    }
118}