1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::layouts::{
  7    AutomorphismKeyToRef, Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  8    prepared::{
  9        GLWESwitchingKeyPrepareTmpBytes, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToMut, GLWESwitchingKeyPreparedToRef,
 10    },
 11};
 12
 13#[derive(PartialEq, Eq)]
 14pub struct AutomorphismKeyPrepared<D: Data, B: Backend> {
 15    pub(crate) key: GLWESwitchingKeyPrepared<D, B>,
 16    pub(crate) p: i64,
 17}
 18
 19impl<D: Data, B: Backend> AutomorphismKeyPrepared<D, B> {
 20    pub fn p(&self) -> i64 {
 21        self.p
 22    }
 23}
 24
 25impl<D: Data, B: Backend> LWEInfos for AutomorphismKeyPrepared<D, B> {
 26    fn n(&self) -> Degree {
 27        self.key.n()
 28    }
 29
 30    fn base2k(&self) -> Base2K {
 31        self.key.base2k()
 32    }
 33
 34    fn k(&self) -> TorusPrecision {
 35        self.key.k()
 36    }
 37
 38    fn size(&self) -> usize {
 39        self.key.size()
 40    }
 41}
 42
 43pub trait SetP {
 44    fn set_p(&mut self, p: i64);
 45}
 46
 47impl<D: Data, B: Backend> SetP for AutomorphismKeyPrepared<D, B> {
 48    fn set_p(&mut self, p: i64) {
 49        self.p = p
 50    }
 51}
 52
 53impl<D: Data, B: Backend> GLWEInfos for AutomorphismKeyPrepared<D, B> {
 54    fn rank(&self) -> Rank {
 55        self.rank_out()
 56    }
 57}
 58
 59impl<D: Data, B: Backend> GGLWEInfos for AutomorphismKeyPrepared<D, B> {
 60    fn rank_in(&self) -> Rank {
 61        self.key.rank_in()
 62    }
 63
 64    fn rank_out(&self) -> Rank {
 65        self.key.rank_out()
 66    }
 67
 68    fn dsize(&self) -> Dsize {
 69        self.key.dsize()
 70    }
 71
 72    fn dnum(&self) -> Dnum {
 73        self.key.dnum()
 74    }
 75}
 76
 77impl<B: Backend> AutomorphismKeyPrepared<Vec<u8>, B> {
 78    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 79    where
 80        A: GGLWEInfos,
 81        Module<B>: VmpPMatAlloc<B>,
 82    {
 83        assert_eq!(
 84            infos.rank_in(),
 85            infos.rank_out(),
 86            "rank_in != rank_out is not supported for GGLWEAutomorphismKeyPrepared"
 87        );
 88        AutomorphismKeyPrepared::<Vec<u8>, B> {
 89            key: GLWESwitchingKeyPrepared::alloc(module, infos),
 90            p: 0,
 91        }
 92    }
 93
 94    pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
 95    where
 96        Module<B>: VmpPMatAlloc<B>,
 97    {
 98        AutomorphismKeyPrepared {
 99            key: GLWESwitchingKeyPrepared::alloc_with(module, base2k, k, rank, rank, dnum, dsize),
100            p: 0,
101        }
102    }
103
104    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
105    where
106        A: GGLWEInfos,
107        Module<B>: VmpPMatAllocBytes,
108    {
109        assert_eq!(
110            infos.rank_in(),
111            infos.rank_out(),
112            "rank_in != rank_out is not supported for GGLWEAutomorphismKeyPrepared"
113        );
114        GLWESwitchingKeyPrepared::alloc_bytes(module, infos)
115    }
116
117    pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
118    where
119        Module<B>: VmpPMatAllocBytes,
120    {
121        GLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rank, rank, dnum, dsize)
122    }
123}
124
125pub trait AutomorphismKeyPrepareTmpBytes {
126    fn automorphism_key_prepare_tmp_bytes<A>(&self, infos: &A)
127    where
128        A: GGLWEInfos;
129}
130
131impl<B: Backend> AutomorphismKeyPrepareTmpBytes for Module<B>
132where
133    Module<B>: GLWESwitchingKeyPrepareTmpBytes,
134{
135    fn automorphism_key_prepare_tmp_bytes<A>(&self, infos: &A)
136    where
137        A: GGLWEInfos,
138    {
139        self.glwe_switching_key_prepare_tmp_bytes(infos)
140    }
141}
142
143impl<D: DataRef, B: Backend> AutomorphismKeyPrepared<D, B> {
144    pub fn prepare_tmp_bytes(&self, module: &Module<B>) -> usize
145    where
146        Module<B>: AutomorphismKeyPrepareTmpBytes,
147    {
148        module.automorphism_key_prepare_tmp_bytes(self);
149    }
150}
151
152pub trait AutomorphismKeyPrepare<B: Backend> {
153    fn automorphism_key_prepare<R, O>(&self, res: &R, other: &O, scratch: &Scratch<B>)
154    where
155        R: AutomorphismKeyPreparedToMut<B>,
156        O: AutomorphismKeyToRef;
157}
158
159impl<B: Backend> AutomorphismKeyPrepare<B> for Module<B> {
160    fn automorphism_key_prepare<R, O>(&self, res: &R, other: &O, scratch: &Scratch<B>)
161    where
162        R: AutomorphismKeyPreparedToMut<B>,
163        O: AutomorphismKeyToRef,
164    {
165        self.key.prepare(self, &other.to_ref().key, scratch);
166        self.p = other.p;
167    }
168}
169
170pub trait AutomorphismKeyPrepareAlloc<B: Backend> {
171    fn automorphism_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> AutomorphismKeyPrepared<Vec<u8>, B>
172    where
173        O: AutomorphismKeyToRef;
174}
175
176impl<B: Backend> AutomorphismKeyPrepareAlloc<B> for Module<B>
177where
178    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
179{
180    fn automorphism_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> AutomorphismKeyPrepared<Vec<u8>, B>
181    where
182        O: AutomorphismKeyToRef,
183    {
184        let mut atk_prepared: AutomorphismKeyPrepared<Vec<u8>, B> = AutomorphismKeyPrepared::alloc(self, &other.to_ref());
185        atk_prepared.prepare(self, &other.to_ref(), scratch);
186        atk_prepared
187    }
188}
189
190pub trait AutomorphismKeyPreparedToMut<B: Backend> {
191    fn to_mut(&mut self) -> AutomorphismKeyPrepared<&mut [u8], B>;
192}
193
194impl<D: DataMut, B: Backend> AutomorphismKeyPreparedToMut<B> for AutomorphismKeyPrepared<D, B> {
195    fn to_mut(&mut self) -> AutomorphismKeyPrepared<&mut [u8], B> {
196        AutomorphismKeyPrepared {
197            p: self.p,
198            key: self.key.to_mut(),
199        }
200    }
201}
202
203pub trait AutomorphismKeyPreparedToRef<B: Backend> {
204    fn to_ref(&self) -> AutomorphismKeyPrepared<&[u8], B>;
205}
206
207impl<D: DataRef, B: Backend> AutomorphismKeyPreparedToRef<B> for AutomorphismKeyPrepared<D, B> {
208    fn to_ref(&self) -> AutomorphismKeyPrepared<&[u8], B> {
209        AutomorphismKeyPrepared {
210            p: self.p,
211            key: self.key.to_ref(),
212        }
213    }
214}