poulpy_core/external_product/
gglwe.rs

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,
34        B: GGSWPreparedToRef<BE>,
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,
66        A: GGLWEToRef,
67        B: GGSWPreparedToRef<BE>,
68        Scratch<BE>: ScratchTakeCore<BE>,
69    {
70        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
71        let a: &GGLWE<&[u8]> = &a.to_ref();
72        let b: &GGSWPrepared<&[u8], BE> = &b.to_ref();
73
74        assert_eq!(
75            res.rank_in(),
76            a.rank_in(),
77            "res input rank_in: {} != a input rank_in: {}",
78            res.rank_in(),
79            a.rank_in()
80        );
81        assert_eq!(
82            a.rank_out(),
83            b.rank(),
84            "a output rank_out: {} != b rank: {}",
85            a.rank_out(),
86            b.rank()
87        );
88        assert_eq!(
89            res.rank_out(),
90            b.rank(),
91            "res output rank_out: {} != b rank: {}",
92            res.rank_out(),
93            b.rank()
94        );
95
96        for row in 0..res.dnum().into() {
97            for col in 0..res.rank_in().into() {
98                self.glwe_external_product(&mut res.at_mut(row, col), &a.at(row, col), b, scratch);
99            }
100        }
101
102        for row in res.dnum().min(a.dnum()).into()..res.dnum().into() {
103            for col in 0..res.rank_in().into() {
104                res.at_mut(row, col).data_mut().zero();
105            }
106        }
107    }
108
109    fn gglwe_external_product_inplace<R, A>(&self, res: &mut R, a: &A, scratch: &mut Scratch<BE>)
110    where
111        R: GGLWEToMut,
112        A: GGSWPreparedToRef<BE>,
113        Scratch<BE>: ScratchTakeCore<BE>,
114    {
115        let res: &mut GGLWE<&mut [u8]> = &mut res.to_mut();
116        let a: &GGSWPrepared<&[u8], BE> = &a.to_ref();
117
118        assert_eq!(
119            res.rank_out(),
120            a.rank(),
121            "res output rank: {} != a rank: {}",
122            res.rank_out(),
123            a.rank()
124        );
125
126        for row in 0..res.dnum().into() {
127            for col in 0..res.rank_in().into() {
128                self.glwe_external_product_inplace(&mut res.at_mut(row, col), a, scratch);
129            }
130        }
131    }
132}
133
134impl<BE: Backend> GGLWEExternalProduct<BE> for Module<BE> where Self: GLWEExternalProduct<BE> {}
135
136impl GLWESwitchingKey<Vec<u8>> {
137    pub fn external_product_tmp_bytes<R, A, B, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, b_infos: &B) -> usize
138    where
139        R: GGLWEInfos,
140        A: GGLWEInfos,
141        B: GGSWInfos,
142        M: GGLWEExternalProduct<BE>,
143    {
144        module.gglwe_external_product_tmp_bytes(res_infos, a_infos, b_infos)
145    }
146}
147
148impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
149    pub fn external_product<A, B, M, BE: Backend>(&mut self, module: &M, a: &A, b: &B, scratch: &mut Scratch<BE>)
150    where
151        M: GGLWEExternalProduct<BE>,
152        A: GGLWEToRef,
153        B: GGSWPreparedToRef<BE>,
154        Scratch<BE>: ScratchTakeCore<BE>,
155    {
156        module.gglwe_external_product(self, a, b, scratch);
157    }
158
159    pub fn external_product_inplace<A, M, BE: Backend>(&mut self, module: &M, a: &A, scratch: &mut Scratch<BE>)
160    where
161        M: GGLWEExternalProduct<BE>,
162        A: GGSWPreparedToRef<BE>,
163        Scratch<BE>: ScratchTakeCore<BE>,
164    {
165        module.gglwe_external_product_inplace(self, a, scratch);
166    }
167}