1use backend::hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat},
  4};
  5
  6use crate::layouts::{GGLWESwitchingKey, Infos, prepared::GGLWECiphertextExec};
  7
  8#[derive(PartialEq, Eq)]
  9pub struct GGLWESwitchingKeyExec<D: Data, B: Backend> {
 10    pub(crate) key: GGLWECiphertextExec<D, B>,
 11    pub(crate) sk_in_n: usize,  // Degree of sk_in
 12    pub(crate) sk_out_n: usize, // Degree of sk_out
 13}
 14
 15impl<B: Backend> GGLWESwitchingKeyExec<Vec<u8>, B> {
 16    pub fn alloc(
 17        module: &Module<B>,
 18        n: usize,
 19        basek: usize,
 20        k: usize,
 21        rows: usize,
 22        digits: usize,
 23        rank_in: usize,
 24        rank_out: usize,
 25    ) -> Self
 26    where
 27        Module<B>: VmpPMatAlloc<B>,
 28    {
 29        GGLWESwitchingKeyExec::<Vec<u8>, B> {
 30            key: GGLWECiphertextExec::alloc(module, n, basek, k, rows, digits, rank_in, rank_out),
 31            sk_in_n: 0,
 32            sk_out_n: 0,
 33        }
 34    }
 35
 36    pub fn bytes_of(
 37        module: &Module<B>,
 38        n: usize,
 39        basek: usize,
 40        k: usize,
 41        rows: usize,
 42        digits: usize,
 43        rank_in: usize,
 44        rank_out: usize,
 45    ) -> usize
 46    where
 47        Module<B>: VmpPMatAllocBytes,
 48    {
 49        GGLWECiphertextExec::bytes_of(module, n, basek, k, rows, digits, rank_in, rank_out)
 50    }
 51
 52    pub fn from<DataOther: DataRef>(module: &Module<B>, other: &GGLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>) -> Self
 53    where
 54        Module<B>: VmpPMatAlloc<B> + VmpPMatPrepare<B>,
 55    {
 56        let mut ksk_exec: GGLWESwitchingKeyExec<Vec<u8>, B> = Self::alloc(
 57            module,
 58            other.n(),
 59            other.basek(),
 60            other.k(),
 61            other.rows(),
 62            other.digits(),
 63            other.rank_in(),
 64            other.rank_out(),
 65        );
 66        ksk_exec.prepare(module, other, scratch);
 67        ksk_exec
 68    }
 69}
 70
 71impl<D: Data, B: Backend> Infos for GGLWESwitchingKeyExec<D, B> {
 72    type Inner = VmpPMat<D, B>;
 73
 74    fn inner(&self) -> &Self::Inner {
 75        self.key.inner()
 76    }
 77
 78    fn basek(&self) -> usize {
 79        self.key.basek()
 80    }
 81
 82    fn k(&self) -> usize {
 83        self.key.k()
 84    }
 85}
 86
 87impl<D: Data, B: Backend> GGLWESwitchingKeyExec<D, B> {
 88    pub fn rank(&self) -> usize {
 89        self.key.data.cols_out() - 1
 90    }
 91
 92    pub fn rank_in(&self) -> usize {
 93        self.key.data.cols_in()
 94    }
 95
 96    pub fn rank_out(&self) -> usize {
 97        self.key.data.cols_out() - 1
 98    }
 99
100    pub fn digits(&self) -> usize {
101        self.key.digits()
102    }
103
104    pub fn sk_degree_in(&self) -> usize {
105        self.sk_in_n
106    }
107
108    pub fn sk_degree_out(&self) -> usize {
109        self.sk_out_n
110    }
111}
112
113impl<D: DataMut, B: Backend> GGLWESwitchingKeyExec<D, B> {
114    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GGLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>)
115    where
116        DataOther: DataRef,
117        Module<B>: VmpPMatPrepare<B>,
118    {
119        self.key.prepare(module, &other.key, scratch);
120        self.sk_in_n = other.sk_in_n;
121        self.sk_out_n = other.sk_out_n;
122    }
123}