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,
  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>(module: &M, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
 99    where
100        R: GGLWEInfos,
101        A: GGLWEInfos,
102        B: GGSWInfos,
103        M: GGLWEExternalProduct<BE>,
104    {
105        module.gglwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
106    }
107}
108
109impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
110    pub fn external_product<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
111    where
112        M: GGLWEExternalProduct<BE>,
113        A: GGLWEToRef,
114        B: GGSWPreparedToRef<BE>,
115        Scratch<BE>: ScratchTakeCore<BE>,
116    {
117        module.gglwe_external_product(self, a, b, scratch);
118    }
119
120    pub fn external_product_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
121    where
122        M: GGLWEExternalProduct<BE>,
123        A: GGSWPreparedToRef<BE>,
124        Scratch<BE>: ScratchTakeCore<BE>,
125    {
126        module.gglwe_external_product_inplace(self, a, scratch);
127    }
128}