1use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
  2
  3use crate::layouts::{
  4    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWEPrepared, GGLWEPreparedFactory, GGLWEPreparedToMut, GGLWEPreparedToRef,
  5    GGLWEToRef, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  6};
  7
  8#[derive(PartialEq, Eq)]
  9pub struct GLWETensorKeyPrepared<D: Data, B: Backend>(pub(crate) GGLWEPrepared<D, B>);
 10
 11impl<D: Data, B: Backend> LWEInfos for GLWETensorKeyPrepared<D, B> {
 12    fn n(&self) -> Degree {
 13        self.0.n()
 14    }
 15
 16    fn base2k(&self) -> Base2K {
 17        self.0.base2k()
 18    }
 19
 20    fn k(&self) -> TorusPrecision {
 21        self.0.k()
 22    }
 23
 24    fn size(&self) -> usize {
 25        self.0.size()
 26    }
 27}
 28
 29impl<D: Data, B: Backend> GLWEInfos for GLWETensorKeyPrepared<D, B> {
 30    fn rank(&self) -> Rank {
 31        self.rank_out()
 32    }
 33}
 34
 35impl<D: Data, B: Backend> GGLWEInfos for GLWETensorKeyPrepared<D, B> {
 36    fn rank_in(&self) -> Rank {
 37        self.0.rank_in()
 38    }
 39
 40    fn rank_out(&self) -> Rank {
 41        self.0.rank_out()
 42    }
 43
 44    fn dsize(&self) -> Dsize {
 45        self.0.dsize()
 46    }
 47
 48    fn dnum(&self) -> Dnum {
 49        self.0.dnum()
 50    }
 51}
 52
 53pub trait GLWETensorKeyPreparedFactory<B: Backend>
 54where
 55    Self: GGLWEPreparedFactory<B>,
 56{
 57    fn alloc_tensor_key_prepared(
 58        &self,
 59        base2k: Base2K,
 60        k: TorusPrecision,
 61        dnum: Dnum,
 62        dsize: Dsize,
 63        rank: Rank,
 64    ) -> GLWETensorKeyPrepared<Vec<u8>, B> {
 65        let pairs: u32 = (((rank.as_u32() + 1) * rank.as_u32()) >> 1).max(1);
 66        GLWETensorKeyPrepared(self.alloc_gglwe_prepared(base2k, k, Rank(pairs), rank, dnum, dsize))
 67    }
 68
 69    fn alloc_tensor_key_prepared_from_infos<A>(&self, infos: &A) -> GLWETensorKeyPrepared<Vec<u8>, B>
 70    where
 71        A: GGLWEInfos,
 72    {
 73        self.alloc_tensor_key_prepared(infos.base2k(), infos.k(), infos.dnum(), infos.dsize(), infos.rank_out())
 74    }
 75
 76    fn bytes_of_tensor_key_prepared(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
 77        let pairs: u32 = (((rank.as_u32() + 1) * rank.as_u32()) >> 1).max(1);
 78        self.bytes_of_gglwe_prepared(base2k, k, Rank(pairs), rank, dnum, dsize)
 79    }
 80
 81    fn bytes_of_tensor_key_prepared_from_infos<A>(&self, infos: &A) -> usize
 82    where
 83        A: GGLWEInfos,
 84    {
 85        self.bytes_of_tensor_key_prepared(infos.base2k(), infos.k(), infos.rank(), infos.dnum(), infos.dsize())
 86    }
 87
 88    fn prepare_tensor_key_tmp_bytes<A>(&self, infos: &A) -> usize
 89    where
 90        A: GGLWEInfos,
 91    {
 92        self.prepare_gglwe_tmp_bytes(infos)
 93    }
 94
 95    fn prepare_tensor_key<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
 96    where
 97        R: GGLWEPreparedToMut<B>,
 98        O: GGLWEToRef,
 99    {
100        self.prepare_gglwe(res, other, scratch);
101    }
102}
103
104impl<B: Backend> GLWETensorKeyPreparedFactory<B> for Module<B> where Module<B>: GGLWEPreparedFactory<B> {}
105
106impl<B: Backend> GLWETensorKeyPrepared<Vec<u8>, B> {
107    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
108    where
109        A: GGLWEInfos,
110        M: GLWETensorKeyPreparedFactory<B>,
111    {
112        module.alloc_tensor_key_prepared_from_infos(infos)
113    }
114
115    pub fn alloc_with<M>(module: &M, base2k: Base2K, k: TorusPrecision, dnum: Dnum, dsize: Dsize, rank: Rank) -> Self
116    where
117        M: GLWETensorKeyPreparedFactory<B>,
118    {
119        module.alloc_tensor_key_prepared(base2k, k, dnum, dsize, rank)
120    }
121
122    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
123    where
124        A: GGLWEInfos,
125        M: GLWETensorKeyPreparedFactory<B>,
126    {
127        module.bytes_of_tensor_key_prepared_from_infos(infos)
128    }
129
130    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
131    where
132        M: GLWETensorKeyPreparedFactory<B>,
133    {
134        module.bytes_of_tensor_key_prepared(base2k, k, rank, dnum, dsize)
135    }
136}
137
138impl<B: Backend> GLWETensorKeyPrepared<Vec<u8>, B> {
139    pub fn prepare_tmp_bytes<A, M>(&self, module: &M, infos: &A) -> usize
140    where
141        A: GGLWEInfos,
142        M: GLWETensorKeyPreparedFactory<B>,
143    {
144        module.prepare_tensor_key_tmp_bytes(infos)
145    }
146}
147
148impl<D: DataMut, B: Backend> GLWETensorKeyPrepared<D, B> {
149    pub fn prepare<O, M>(&mut self, module: &M, other: &O, scratch: &mut Scratch<B>)
150    where
151        O: GGLWEToRef,
152        M: GLWETensorKeyPreparedFactory<B>,
153    {
154        module.prepare_tensor_key(self, other, scratch);
155    }
156}
157
158impl<D: DataMut, B: Backend> GGLWEPreparedToMut<B> for GLWETensorKeyPrepared<D, B>
159where
160    GGLWEPrepared<D, B>: GGLWEPreparedToMut<B>,
161{
162    fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B> {
163        self.0.to_mut()
164    }
165}
166
167impl<D: DataRef, B: Backend> GGLWEPreparedToRef<B> for GLWETensorKeyPrepared<D, B>
168where
169    GGLWEPrepared<D, B>: GGLWEPreparedToRef<B>,
170{
171    fn to_ref(&self) -> GGLWEPrepared<&[u8], B> {
172        self.0.to_ref()
173    }
174}