1use poulpy_hal::{
  2    api::{ModuleN, ScratchAvailable},
  3    layouts::{Backend, DataMut, Module, Scratch, ZnxZero},
  4};
  5
  6use crate::{
  7    GLWEExternalProduct, ScratchTakeCore,
  8    layouts::{
  9        GGSW, GGSWInfos, GGSWToMut, GGSWToRef, GLWEInfos, LWEInfos,
 10        prepared::{GGSWPrepared, GGSWPreparedToRef},
 11    },
 12};
 13
 14pub trait GGSWExternalProduct<BE: Backend>
 15where
 16    Self: GLWEExternalProduct<BE> + ModuleN,
 17{
 18    fn ggsw_external_product_tmp_bytes<R, A, B>(&self, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
 19    where
 20        R: GGSWInfos,
 21        A: GGSWInfos,
 22        B: GGSWInfos,
 23    {
 24        self.glwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
 25    }
 26
 27    fn ggsw_external_product<R, A, B>(&self, res: &mut R, a: &A, b: &B, scratch: &mut Scratch<BE>)
 28    where
 29        R: GGSWToMut,
 30        A: GGSWToRef,
 31        B: GGSWPreparedToRef<BE>,
 32        Scratch<BE>: ScratchTakeCore<BE>,
 33    {
 34        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 35        let a: &GGSW<&[u8]> = &a.to_ref();
 36        let b: &GGSWPrepared<&[u8], BE> = &b.to_ref();
 37
 38        assert_eq!(res.rank(), a.rank(), "res rank: {} != a rank: {}", res.rank(), a.rank());
 39        assert_eq!(res.rank(), b.rank(), "res rank: {} != b rank: {}", res.rank(), b.rank());
 40
 41        assert_eq!(res.base2k(), a.base2k());
 42
 43        assert!(scratch.available() >= self.ggsw_external_product_tmp_bytes(res, a, b));
 44
 45        let min_dnum: usize = res.dnum().min(a.dnum()).into();
 46
 47        for row in 0..min_dnum {
 48            for col in 0..(res.rank() + 1).into() {
 49                self.glwe_external_product(&mut res.at_mut(row, col), &a.at(row, col), b, scratch);
 50            }
 51        }
 52
 53        for row in min_dnum..res.dnum().into() {
 54            for col in 0..(res.rank() + 1).into() {
 55                res.at_mut(row, col).data.zero();
 56            }
 57        }
 58    }
 59
 60    fn ggsw_external_product_inplace<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
 61    where
 62        R: GGSWToMut,
 63        A: GGSWPreparedToRef<BE>,
 64        Scratch<BE>: ScratchTakeCore<BE>,
 65    {
 66        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
 67        let a: &GGSWPrepared<&[u8], BE> = &a.to_ref();
 68
 69        assert_eq!(res.n(), self.n() as u32);
 70        assert_eq!(a.n(), self.n() as u32);
 71        assert_eq!(res.rank(), a.rank(), "res rank: {} != a rank: {}", res.rank(), a.rank());
 72
 73        for row in 0..res.dnum().into() {
 74            for col in 0..(res.rank() + 1).into() {
 75                self.glwe_external_product_inplace(&mut res.at_mut(row, col), a, scratch);
 76            }
 77        }
 78    }
 79}
 80
 81impl<BE: Backend> GGSWExternalProduct<BE> for Module<BE> where Self: GLWEExternalProduct<BE> {}
 82
 83impl GGSW<Vec<u8>> {
 84    pub fn external_product_tmp_bytes<R, A, B, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
 85    where
 86        R: GGSWInfos,
 87        A: GGSWInfos,
 88        B: GGSWInfos,
 89        M: GGSWExternalProduct<BE>,
 90    {
 91        module.ggsw_external_product_tmp_bytes(res_infos, a_infos, b_infos)
 92    }
 93}
 94
 95impl<DataSelf: DataMut> GGSW<DataSelf> {
 96    pub fn external_product<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
 97    where
 98        M: GGSWExternalProduct<BE>,
 99        A: GGSWToRef,
100        B: GGSWPreparedToRef<BE>,
101        Scratch<BE>: ScratchTakeCore<BE>,
102    {
103        module.ggsw_external_product(self, a, b, scratch);
104    }
105
106    pub fn external_product_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
107    where
108        M: GGSWExternalProduct<BE>,
109        A: GGSWPreparedToRef<BE>,
110        Scratch<BE>: ScratchTakeCore<BE>,
111    {
112        module.ggsw_external_product_inplace(self, a, scratch);
113    }
114}