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