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