poulpy/spqlios/src/module.rs
spqlios/src/module.rsRaw
1use crate::bindings::*;
2use crate::poly::Poly;
3use crate::scalar::Scalar;
4
5pub type MODULETYPE = u8;
6pub const FFT64: u8 = 0;
7pub const NTT120: u8 = 1;
8
9pub struct Module(*mut MODULE);
10
11impl Module {
12 // Instantiates a new module.
13 pub fn new<const MODULETYPE: MODULETYPE>(n: usize) -> Self {
14 unsafe {
15 let m: *mut module_info_t = new_module_info(n as u64, MODULETYPE as u32);
16 if m.is_null() {
17 panic!("Failed to create module.");
18 }
19 Self(m)
20 }
21 }
22
23 // Prepares a scalar polynomial (1 limb) for a scalar x vector product.
24 // Method will panic if a.limbs() != 1.
25 pub fn svp_prepare(&self, svp_ppol: &mut SVPPOL, a: &Scalar) {
26 unsafe { svp_prepare(self.0, svp_ppol.0, a.as_ptr()) }
27 }
28
29 // Allocates a scalar-vector-product prepared-poly (SVPPOL).
30 pub fn new_svp_ppol(&self) -> SVPPOL {
31 unsafe { SVPPOL(new_svp_ppol(self.0)) }
32 }
33
34 // Allocates a vector Z[X]/(X^N+1) that stores normalized in the DFT space.
35 pub fn new_vec_znx_dft(&self, limbs: usize) -> VECZNXDFT {
36 unsafe { VECZNXDFT(new_vec_znx_dft(self.0, limbs as u64), limbs) }
37 }
38
39 // Allocates a vector Z[X]/(X^N+1) that stores not normalized values.
40 pub fn new_vec_znx_big(&self, limbs: usize) -> VECZNXBIG {
41 unsafe { VECZNXBIG(new_vec_znx_big(self.0, limbs as u64), limbs) }
42 }
43
44 // Applies a scalar x vector product: res <- a (ppol) x b
45 pub fn svp_apply_dft(&self, c: &mut VECZNXDFT, a: &SVPPOL, b: &Poly) {
46 let limbs: u64 = b.limbs() as u64;
47 assert!(
48 c.limbs() as u64 >= limbs,
49 "invalid c_vector: c_vector.limbs()={} < b.limbs()={}",
50 c.limbs(),
51 limbs
52 );
53 unsafe { svp_apply_dft(self.0, c.0, limbs, a.0, b.as_ptr(), limbs, b.n() as u64) }
54 }
55
56 // b <- IDFT(a), uses a as scratch space.
57 pub fn vec_znx_idft_tmp_a(&self, b: &mut VECZNXBIG, a: &mut VECZNXDFT, a_limbs: usize) {
58 assert!(
59 b.limbs() >= a_limbs,
60 "invalid c_vector: b_vector.limbs()={} < a_limbs={}",
61 b.limbs(),
62 a_limbs
63 );
64 unsafe { vec_znx_idft_tmp_a(self.0, b.0, a_limbs as u64, a.0, a_limbs as u64) }
65 }
66
67 // Returns the size of the scratch space for [vec_znx_idft].
68 pub fn vec_znx_idft_tmp_bytes(&self) -> usize {
69 unsafe { vec_znx_idft_tmp_bytes(self.0) as usize }
70 }
71
72 // b <- IDFT(a), scratch space size obtained with [vec_znx_idft_tmp_bytes].
73 pub fn vec_znx_idft(
74 &self,
75 b_vector: &mut VECZNXBIG,
76 a_vector: &mut VECZNXDFT,
77 a_limbs: usize,
78 tmp_bytes: &mut [u8],
79 ) {
80 assert!(
81 b_vector.limbs() >= a_limbs,
82 "invalid c_vector: b_vector.limbs()={} < a_limbs={}",
83 b_vector.limbs(),
84 a_limbs
85 );
86 assert!(
87 a_vector.limbs() >= a_limbs,
88 "invalid c_vector: c_vector.limbs()={} < a_limbs={}",
89 a_vector.limbs(),
90 a_limbs
91 );
92 assert!(
93 tmp_bytes.len() <= self.vec_znx_idft_tmp_bytes(),
94 "invalid tmp_bytes: tmp_bytes.len()={} < self.vec_znx_idft_tmp_bytes()={}",
95 tmp_bytes.len(),
96 self.vec_znx_idft_tmp_bytes()
97 );
98 unsafe {
99 vec_znx_idft(
100 self.0,
101 b_vector.0,
102 a_limbs as u64,
103 a_vector.0,
104 a_limbs as u64,
105 tmp_bytes.as_mut_ptr(),
106 )
107 }
108 }
109
110 // c <- b - a
111 pub fn vec_znx_big_sub_small_a(&self, c: &mut VECZNXBIG, a: &Poly, b: &VECZNXBIG) {
112 let limbs: usize = a.limbs();
113 assert!(
114 b.limbs() >= limbs,
115 "invalid c: b.limbs()={} < a.limbs()={}",
116 b.limbs(),
117 limbs
118 );
119 assert!(
120 c.limbs() >= limbs,
121 "invalid c: c.limbs()={} < a.limbs()={}",
122 c.limbs(),
123 limbs
124 );
125 unsafe {
126 vec_znx_big_sub_small_a(
127 self.0,
128 c.0,
129 c.limbs() as u64,
130 a.as_ptr(),
131 limbs as u64,
132 a.n() as u64,
133 b.0,
134 b.limbs() as u64,
135 )
136 }
137 }
138
139 // b <- b - a
140 pub fn vec_znx_big_sub_small_a_inplace(&self, b: &mut VECZNXBIG, a: &Poly) {
141 let limbs: usize = a.limbs();
142 assert!(
143 b.limbs() >= limbs,
144 "invalid c_vector: b.limbs()={} < a.limbs()={}",
145 b.limbs(),
146 limbs
147 );
148 unsafe {
149 vec_znx_big_sub_small_a(
150 self.0,
151 b.0,
152 b.limbs() as u64,
153 a.as_ptr(),
154 limbs as u64,
155 a.n() as u64,
156 b.0,
157 b.limbs() as u64,
158 )
159 }
160 }
161
162 // c <- b + a
163 pub fn vec_znx_big_add_small(&self, c: &mut VECZNXBIG, a: &Poly, b: &VECZNXBIG) {
164 let limbs: usize = a.limbs();
165 assert!(
166 b.limbs() >= limbs,
167 "invalid c: b.limbs()={} < a.limbs()={}",
168 b.limbs(),
169 limbs
170 );
171 assert!(
172 c.limbs() >= limbs,
173 "invalid c: c.limbs()={} < a.limbs()={}",
174 c.limbs(),
175 limbs
176 );
177 unsafe {
178 vec_znx_big_add_small(
179 self.0,
180 c.0,
181 limbs as u64,
182 b.0,
183 limbs as u64,
184 a.as_ptr(),
185 limbs as u64,
186 a.n() as u64,
187 )
188 }
189 }
190
191 // b <- b + a
192 pub fn vec_znx_big_add_small_inplace(&self, b: &mut VECZNXBIG, a: &Poly) {
193 let limbs: usize = a.limbs();
194 assert!(
195 b.limbs() >= limbs,
196 "invalid c_vector: b.limbs()={} < a.limbs()={}",
197 b.limbs(),
198 limbs
199 );
200 unsafe {
201 vec_znx_big_add_small(
202 self.0,
203 b.0,
204 limbs as u64,
205 b.0,
206 limbs as u64,
207 a.as_ptr(),
208 limbs as u64,
209 a.n() as u64,
210 )
211 }
212 }
213
214 pub fn vec_znx_big_normalize_tmp_bytes(&self) -> usize {
215 unsafe { vec_znx_big_normalize_base2k_tmp_bytes(self.0) as usize }
216 }
217
218 // b <- normalize(a)
219 pub fn vec_znx_big_normalize(&self, b: &mut Poly, a: &VECZNXBIG, tmp_bytes: &mut [u8]) {
220 let limbs: usize = b.limbs();
221 assert!(
222 b.limbs() >= limbs,
223 "invalid c_vector: b.limbs()={} < a.limbs()={}",
224 b.limbs(),
225 limbs
226 );
227 assert!(
228 tmp_bytes.len() <= self.vec_znx_big_normalize_tmp_bytes(),
229 "invalid tmp_bytes: tmp_bytes.len()={} <= self.vec_znx_big_normalize_tmp_bytes()={}",
230 tmp_bytes.len(),
231 self.vec_znx_big_normalize_tmp_bytes()
232 );
233 unsafe {
234 vec_znx_big_normalize_base2k(
235 self.0,
236 b.log_base2k as u64,
237 b.as_mut_ptr(),
238 limbs as u64,
239 b.n() as u64,
240 a.0,
241 limbs as u64,
242 tmp_bytes.as_mut_ptr(),
243 )
244 }
245 }
246}
247
248pub struct SVPPOL(*mut svp_ppol_t);
249pub struct VECZNXDFT(*mut vec_znx_dft_t, usize);
250pub struct VECZNXBIG(*mut vec_znx_bigcoeff_t, usize);
251
252impl VECZNXBIG {
253 pub fn as_vec_znx_dft(&mut self) -> VECZNXDFT {
254 VECZNXDFT(self.0 as *mut vec_znx_dft_t, self.1)
255 }
256 pub fn limbs(&self) -> usize {
257 self.1
258 }
259}
260
261impl VECZNXDFT {
262 pub fn as_vec_znx_big(&mut self) -> VECZNXBIG {
263 VECZNXBIG(self.0 as *mut vec_znx_bigcoeff_t, self.1)
264 }
265 pub fn limbs(&self) -> usize {
266 self.1
267 }
268}