Skip to main content

poulpy_hal/api/
vec_znx_dft.rs

1use crate::layouts::{
2    Backend, ScratchArena, VecZnxBackendRef, VecZnxBigBackendMut, VecZnxDftBackendMut, VecZnxDftBackendRef, VecZnxDftOwned,
3};
4
5/// Allocates a [`VecZnxDft`](crate::layouts::VecZnxDft).
6pub trait VecZnxDftAlloc<B: Backend> {
7    fn vec_znx_dft_alloc(&self, cols: usize, size: usize) -> VecZnxDftOwned<B>;
8}
9
10/// Wraps a byte buffer into a [`VecZnxDft`](crate::layouts::VecZnxDft).
11pub trait VecZnxDftFromBytes<B: Backend> {
12    fn vec_znx_dft_from_bytes(&self, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<B>;
13}
14
15/// Returns the byte size required for a [`VecZnxDft`](crate::layouts::VecZnxDft).
16pub trait VecZnxDftBytesOf {
17    fn bytes_of_vec_znx_dft(&self, cols: usize, size: usize) -> usize;
18}
19
20/// Applies the forward DFT to a coefficient-domain [`VecZnx`](crate::layouts::VecZnx),
21/// storing the result in a [`VecZnxDft`](crate::layouts::VecZnxDft).
22///
23/// The `step` and `offset` parameters select which limbs of the input
24/// are transformed: limbs `offset, offset + step, offset + 2*step, ...`.
25pub trait VecZnxDftApply<B: Backend> {
26    fn vec_znx_dft_apply(
27        &self,
28        step: usize,
29        offset: usize,
30        res: &mut VecZnxDftBackendMut<'_, B>,
31        res_col: usize,
32        a: &VecZnxBackendRef<'_, B>,
33        a_col: usize,
34    );
35}
36
37/// Returns scratch bytes required for [`VecZnxIdftApply`].
38pub trait VecZnxIdftApplyTmpBytes {
39    fn vec_znx_idft_apply_tmp_bytes(&self) -> usize;
40}
41
42/// Applies the inverse DFT, converting a [`VecZnxDft`](crate::layouts::VecZnxDft)
43/// into a [`VecZnxBig`](crate::layouts::VecZnxBig) (extended precision).
44pub trait VecZnxIdftApply<B: Backend> {
45    fn vec_znx_idft_apply(
46        &self,
47        res: &mut VecZnxBigBackendMut<'_, B>,
48        res_col: usize,
49        a: &VecZnxDftBackendRef<'_, B>,
50        a_col: usize,
51        scratch: &mut ScratchArena<'_, B>,
52    );
53}
54
55/// Inverse DFT using `a` as temporary storage (avoids extra scratch).
56pub trait VecZnxIdftApplyTmpA<B: Backend> {
57    fn vec_znx_idft_apply_tmpa(
58        &self,
59        res: &mut VecZnxBigBackendMut<'_, B>,
60        res_col: usize,
61        a: &mut VecZnxDftBackendMut<'_, B>,
62        a_col: usize,
63    );
64}
65
66/// Element-wise addition of two [`VecZnxDft`](crate::layouts::VecZnxDft) vectors.
67pub trait VecZnxDftAddInto<B: Backend> {
68    fn vec_znx_dft_add_into(
69        &self,
70        res: &mut VecZnxDftBackendMut<'_, B>,
71        res_col: usize,
72        a: &VecZnxDftBackendRef<'_, B>,
73        a_col: usize,
74        b: &VecZnxDftBackendRef<'_, B>,
75        b_col: usize,
76    );
77}
78
79/// In-place addition in DFT domain: `res += a`.
80pub trait VecZnxDftAddAssign<B: Backend> {
81    fn vec_znx_dft_add_assign(
82        &self,
83        res: &mut VecZnxDftBackendMut<'_, B>,
84        res_col: usize,
85        a: &VecZnxDftBackendRef<'_, B>,
86        a_col: usize,
87    );
88}
89
90/// In-place LIMB-SHIFTED addition in the DFT domain:
91/// `res += a * 2^(a_scale * base2k)` (`a_scale` is a limb offset, positive
92/// shifts toward the most significant limb — NOT an integer scaling of `a`).
93pub trait VecZnxDftAddScaledAssign<B: Backend> {
94    fn vec_znx_dft_add_scaled_assign(
95        &self,
96        res: &mut VecZnxDftBackendMut<'_, B>,
97        res_col: usize,
98        a: &VecZnxDftBackendRef<'_, B>,
99        a_col: usize,
100        a_scale: i64,
101    );
102}
103
104/// Element-wise subtraction of two [`VecZnxDft`](crate::layouts::VecZnxDft) vectors.
105pub trait VecZnxDftSub<B: Backend> {
106    fn vec_znx_dft_sub(
107        &self,
108        res: &mut VecZnxDftBackendMut<'_, B>,
109        res_col: usize,
110        a: &VecZnxDftBackendRef<'_, B>,
111        a_col: usize,
112        b: &VecZnxDftBackendRef<'_, B>,
113        b_col: usize,
114    );
115}
116
117/// In-place subtraction in DFT domain: `res -= a`.
118pub trait VecZnxDftSubAssign<B: Backend> {
119    fn vec_znx_dft_sub_assign(
120        &self,
121        res: &mut VecZnxDftBackendMut<'_, B>,
122        res_col: usize,
123        a: &VecZnxDftBackendRef<'_, B>,
124        a_col: usize,
125    );
126}
127
128/// In-place negated subtraction in DFT domain: `res = a - res`.
129pub trait VecZnxDftSubNegateAssign<B: Backend> {
130    fn vec_znx_dft_sub_negate_assign(
131        &self,
132        res: &mut VecZnxDftBackendMut<'_, B>,
133        res_col: usize,
134        a: &VecZnxDftBackendRef<'_, B>,
135        a_col: usize,
136    );
137}
138
139/// Copies selected limbs from one [`VecZnxDft`](crate::layouts::VecZnxDft) to another.
140///
141/// The `step` and `offset` parameters select which limbs are copied.
142pub trait VecZnxDftCopy<B: Backend> {
143    fn vec_znx_dft_copy(
144        &self,
145        step: usize,
146        offset: usize,
147        res: &mut VecZnxDftBackendMut<'_, B>,
148        res_col: usize,
149        a: &VecZnxDftBackendRef<'_, B>,
150        a_col: usize,
151    );
152}
153
154/// Zeroes all limbs of the selected column in DFT domain.
155pub trait VecZnxDftZero<B: Backend> {
156    fn vec_znx_dft_zero(&self, res: &mut VecZnxDftBackendMut<'_, B>, res_col: usize);
157}
158
159/// Builds a backend-specific permutation plan that implements the DFT-domain
160/// automorphism `tau_p: X -> X^p` for odd `p`. The plan captures the
161/// slot↔slot permutation (plus any backend-specific bookkeeping such as a
162/// half-spectrum conjugate flag) and is reusable across columns and limbs.
163///
164/// The associated `Plan` type is the only point in the public API where
165/// the backend leaks its automorphism representation. Callers that want to
166/// keep plans backend-agnostic must own
167/// `<Module<B> as VecZnxDftAutomorphismPlan<B>>::Plan`.
168pub trait VecZnxDftAutomorphismPlan<B: Backend> {
169    type Plan;
170
171    fn vec_znx_dft_automorphism_plan(&self, p: i64) -> Self::Plan;
172}
173
174/// Applies a precomputed DFT-domain automorphism plan to `a`, writing the
175/// result into `res` (out-of-place).
176pub trait VecZnxDftAutomorphism<B: Backend>: VecZnxDftAutomorphismPlan<B> {
177    fn vec_znx_dft_automorphism_with_plan(
178        &self,
179        plan: &Self::Plan,
180        res: &mut VecZnxDftBackendMut<'_, B>,
181        res_col: usize,
182        a: &VecZnxDftBackendRef<'_, B>,
183        a_col: usize,
184    );
185
186    /// Convenience: build the plan and apply in one call. Prefer
187    /// [`vec_znx_dft_automorphism_with_plan`](Self::vec_znx_dft_automorphism_with_plan)
188    /// when the same `p` is used repeatedly.
189    fn vec_znx_dft_automorphism(
190        &self,
191        p: i64,
192        res: &mut VecZnxDftBackendMut<'_, B>,
193        res_col: usize,
194        a: &VecZnxDftBackendRef<'_, B>,
195        a_col: usize,
196    ) {
197        let plan = self.vec_znx_dft_automorphism_plan(p);
198        self.vec_znx_dft_automorphism_with_plan(&plan, res, res_col, a, a_col);
199    }
200}