poulpy_hal/api/
vec_znx_dft.rs

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

poulpy_hal/api/
vec_znx_dft.rs
