1use poulpy_hal::layouts::{Backend, DataMut, Module, Scratch, ZnxZero};
  2
  3use crate::{
  4    GLWEExternalProduct, ScratchTakeCore,
  5    layouts::{
  6        GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GGSWInfos, GGSWPrepared, GLWEAutomorphismKey, GLWEInfos, GLWESwitchingKey,
  7        prepared::GGSWPreparedToRef,
  8    },
  9};
 10
 11impl GLWEAutomorphismKey<Vec<u8>> {
 12    pub fn external_product_tmp_bytes<R, A, B, M, BE: Backend>(
 13        &self,
 14        module: &M,
 15        res_infos: &R,
 16        a_infos: &A,
 17        b_infos: &B,
 18    ) -> usize
 19    where
 20        R: GGLWEInfos,
 21        A: GGLWEInfos,
 22        B: GGSWInfos,
 23        M: GGLWEExternalProduct<BE>,
 24    {
 25        module.gglwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
 26    }
 27}
 28
 29impl<DataSelf: DataMut> GLWEAutomorphismKey<DataSelf> {
 30    pub fn external_product<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
 31    where
 32        M: GGLWEExternalProduct<BE>,
 33        A: GGLWEToRef + GGLWEInfos,
 34        B: GGSWPreparedToRef<BE> + GGSWInfos,
 35        Scratch<BE>: ScratchTakeCore<BE>,
 36    {
 37        module.gglwe_external_product(self, a, b, scratch);
 38    }
 39
 40    pub fn external_product_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
 41    where
 42        M: GGLWEExternalProduct<BE>,
 43        A: GGSWPreparedToRef<BE>,
 44        Scratch<BE>: ScratchTakeCore<BE>,
 45    {
 46        module.gglwe_external_product_inplace(self, a, scratch);
 47    }
 48}
 49
 50pub trait GGLWEExternalProduct<BE: Backend>
 51where
 52    Self: GLWEExternalProduct<BE>,
 53{
 54    fn gglwe_external_product_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
 55    where
 56        R: GGLWEInfos,
 57        A: GGLWEInfos,
 58        B: GGSWInfos,
 59    {
 60        self.glwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
 61    }
 62
 63    fn gglwe_external_product<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch<BE>)
 64    where
 65        R: GGLWEToMut + GGLWEInfos,
 66        A: GGLWEToRef + GGLWEInfos,
 67        B: GGSWPreparedToRef<BE> + GGSWInfos,
 68        Scratch<BE>: ScratchTakeCore<BE>,
 69    {
 70        assert_eq!(
 71            res.rank_in(),
 72            a.rank_in(),
 73            "res input rank_in: {} != a input rank_in: {}",
 74            res.rank_in(),
 75            a.rank_in()
 76        );
 77        assert_eq!(
 78            a.rank_out(),
 79            b.rank(),
 80            "a output rank_out: {} != b rank: {}",
 81            a.rank_out(),
 82            b.rank()
 83        );
 84        assert_eq!(
 85            res.rank_out(),
 86            b.rank(),
 87            "res output rank_out: {} != b rank: {}",
 88            res.rank_out(),
 89            b.rank()
 90        );
 91        assert_eq!(res.base2k(), a.base2k());
 92
 93        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
 94        let a: &GGLWE<&[u8]> = &a.to_ref();
 95        let b: &GGSWPrepared<&[u8], BE> = &b.to_ref();
 96
 97        for row in 0..res.dnum().into() {
 98            for col in 0..res.rank_in().into() {
 99                self.glwe_external_product(&mut res.at_mut(row, col), &a.at(row, col), b, scratch);
100            }
101        }
102
103        for row in res.dnum().min(a.dnum()).into()..res.dnum().into() {
104            for col in 0..res.rank_in().into() {
105                res.at_mut(row, col).data_mut().zero();
106            }
107        }
108    }
109
110    fn gglwe_external_product_inplace<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
111    where
112        R: GGLWEToMut,
113        A: GGSWPreparedToRef<BE>,
114        Scratch<BE>: ScratchTakeCore<BE>,
115    {
116        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
117        let a: &GGSWPrepared<&[u8], BE> = &a.to_ref();
118
119        assert_eq!(
120            res.rank_out(),
121            a.rank(),
122            "res output rank: {} != a rank: {}",
123            res.rank_out(),
124            a.rank()
125        );
126
127        for row in 0..res.dnum().into() {
128            for col in 0..res.rank_in().into() {
129                self.glwe_external_product_inplace(&mut res.at_mut(row, col), a, scratch);
130            }
131        }
132    }
133}
134
135impl<BE: Backend> GGLWEExternalProduct<BE> for Module<BE> where Self: GLWEExternalProduct<BE> {}
136
137impl GLWESwitchingKey<Vec<u8>> {
138    pub fn external_product_tmp_bytes<R, A, B, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
139    where
140        R: GGLWEInfos,
141        A: GGLWEInfos,
142        B: GGSWInfos,
143        M: GGLWEExternalProduct<BE>,
144    {
145        module.gglwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
146    }
147}
148
149impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
150    pub fn external_product<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
151    where
152        M: GGLWEExternalProduct<BE>,
153        A: GGLWEToRef + GGLWEInfos,
154        B: GGSWPreparedToRef<BE> + GGSWInfos,
155        Scratch<BE>: ScratchTakeCore<BE>,
156    {
157        module.gglwe_external_product(self, a, b, scratch);
158    }
159
160    pub fn external_product_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
161    where
162        M: GGLWEExternalProduct<BE>,
163        A: GGSWPreparedToRef<BE>,
164        Scratch<BE>: ScratchTakeCore<BE>,
165    {
166        module.gglwe_external_product_inplace(self, a, scratch);
167    }
168}