poulpy/math/src/ring/impl_u64/automorphism.rs
math/src/ring/impl_u64/automorphism.rsRaw
1use crate::modulus::{ScalarOperations, ONCE};
2use crate::modulus::{WordOps, REDUCEMOD};
3use crate::poly::Poly;
4use crate::ring::Ring;
5
6/// Returns a lookup table for the automorphism X^{i} -> X^{i * k mod nth_root}.
7/// Method will panic if n or nth_root are not power-of-two.
8/// Method will panic if gal_el is not coprime with nth_root.
9pub fn automorphism_index<const NTT: bool>(n: usize, nth_root: usize, gal_el: usize) -> Vec<usize> {
10 assert!(n & (n - 1) != 0, "invalid n={}: not a power-of-two", n);
11 assert!(
12 nth_root & (nth_root - 1) == 0,
13 "invalid nth_root={}: not a power-of-two",
14 n
15 );
16 assert!(
17 gal_el & 1 == 1,
18 "invalid gal_el={}: not coprime with nth_root={}",
19 gal_el,
20 nth_root
21 );
22
23 let mut index: Vec<usize> = Vec::with_capacity(n);
24
25 if NTT {
26 let mask = nth_root - 1;
27 let log_nth_root_half: u32 = nth_root.log2() as u32 - 1;
28 for i in 0..n {
29 let i_rev: usize = 2 * i.reverse_bits_msb(log_nth_root_half) + 1;
30 let gal_el_i: usize = ((gal_el * i_rev) & mask) >> 1;
31 index.push(gal_el_i.reverse_bits_msb(log_nth_root_half));
32 }
33 } else {
34 let log_n: usize = n.log2();
35 let mask: usize = (n - 1) as usize;
36 for i in 0..n {
37 let gal_el_i: usize = i as usize * gal_el;
38 let sign: usize = (gal_el_i >> log_n) & 1;
39 let i_out: usize = (gal_el_i & mask) | (sign << (usize::BITS - 1));
40 index.push(i_out)
41 }
42 }
43
44 index
45}
46
47impl Ring<u64> {
48 // b <- auto(a)
49 pub fn a_apply_automorphism_into_b<const NTT: bool>(
50 &self,
51 a: &Poly<u64>,
52 gal_el: usize,
53 nth_root: usize,
54 b: &mut Poly<u64>,
55 ) {
56 self.apply_automorphism_core::<0, ONCE, NTT>(a, gal_el, nth_root, b)
57 }
58
59 // b <- REDUCEMOD(b + auto(a))
60 pub fn a_apply_automorphism_add_b_into_b<const REDUCE: REDUCEMOD, const NTT: bool>(
61 &self,
62 a: &Poly<u64>,
63 gal_el: usize,
64 nth_root: usize,
65 b: &mut Poly<u64>,
66 ) {
67 self.apply_automorphism_core::<1, REDUCE, NTT>(a, gal_el, nth_root, b)
68 }
69
70 // b <- REDUCEMOD(b - auto(a))
71 pub fn a_apply_automorphism_sub_b_into_b<const REDUCE: REDUCEMOD, const NTT: bool>(
72 &self,
73 a: &Poly<u64>,
74 gal_el: usize,
75 nth_root: usize,
76 b: &mut Poly<u64>,
77 ) {
78 self.apply_automorphism_core::<2, REDUCE, NTT>(a, gal_el, nth_root, b)
79 }
80
81 fn apply_automorphism_core<const MOD: u8, const REDUCE: REDUCEMOD, const NTT: bool>(
82 &self,
83 a: &Poly<u64>,
84 gal_el: usize,
85 nth_root: usize,
86 b: &mut Poly<u64>,
87 ) {
88 debug_assert!(
89 a.n() == b.n(),
90 "invalid inputs: a.n() = {} != b.n() = {}",
91 a.n(),
92 b.n()
93 );
94
95 assert!(
96 gal_el & 1 == 1,
97 "invalid gal_el={}: not coprime with nth_root={}",
98 gal_el,
99 nth_root
100 );
101
102 assert!(
103 nth_root & (nth_root - 1) == 0,
104 "invalid nth_root={}: not a power-of-two",
105 nth_root
106 );
107
108 let b_vec: &mut Vec<u64> = &mut b.0;
109 let a_vec: &Vec<u64> = &a.0;
110
111 if NTT {
112 let mask: usize = nth_root - 1;
113 let log_nth_root_half: u32 = nth_root.log2() as u32 - 1;
114 a_vec.iter().enumerate().for_each(|(i, ai)| {
115 let i_rev: usize = 2 * i.reverse_bits_msb(log_nth_root_half) + 1;
116 let gal_el_i: usize = (((gal_el * i_rev) & mask) - 1) >> 1;
117 let idx: usize = gal_el_i.reverse_bits_msb(log_nth_root_half);
118 match MOD {
119 0 => b_vec[idx] = *ai,
120 1 => self
121 .modulus
122 .sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx]),
123 2 => self
124 .modulus
125 .sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx]),
126 _ => {
127 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
128 }
129 }
130 });
131 } else {
132 let n: usize = a.n();
133 let mask: usize = n - 1;
134 let log_n: usize = n.log2();
135 let q: u64 = self.modulus.q();
136 a_vec.iter().enumerate().for_each(|(i, ai)| {
137 let gal_el_i: usize = i * gal_el;
138 let sign: u64 = ((gal_el_i >> log_n) & 1) as u64;
139 let i_out: usize = gal_el_i & mask;
140 let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
141 match MOD {
142 0 => b_vec[i_out] = v,
143 1 => self
144 .modulus
145 .sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[i_out]),
146 2 => self
147 .modulus
148 .sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[i_out]),
149 _ => {
150 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
151 }
152 }
153 });
154 }
155 }
156
157 // b <- auto(a)
158 pub fn a_apply_automorphism_from_index_into_b<const NTT: bool>(
159 &self,
160 a: &Poly<u64>,
161 idx: &[usize],
162 b: &mut Poly<u64>,
163 ) {
164 self.automorphism_from_index_core::<0, ONCE, NTT>(a, idx, b)
165 }
166
167 // b <- REDUCEMOD(b + auto(a))
168 pub fn a_apply_automorphism_from_index_add_b_into_b<
169 const REDUCE: REDUCEMOD,
170 const NTT: bool,
171 >(
172 &self,
173 a: &Poly<u64>,
174 idx: &[usize],
175 b: &mut Poly<u64>,
176 ) {
177 self.automorphism_from_index_core::<1, REDUCE, NTT>(a, idx, b)
178 }
179
180 // b <- REDUCEMOD(b - auto(a))
181 pub fn a_apply_automorphism_from_index_sub_b_into_b<
182 const REDUCE: REDUCEMOD,
183 const NTT: bool,
184 >(
185 &self,
186 a: &Poly<u64>,
187 idx: &[usize],
188 b: &mut Poly<u64>,
189 ) {
190 self.automorphism_from_index_core::<2, REDUCE, NTT>(a, idx, b)
191 }
192
193 // b <- auto(a) if OVERWRITE else b <- REDUCEMOD(b + auto(a))
194 fn automorphism_from_index_core<const MOD: u8, const REDUCE: REDUCEMOD, const NTT: bool>(
195 &self,
196 a: &Poly<u64>,
197 idx: &[usize],
198 b: &mut Poly<u64>,
199 ) {
200 debug_assert!(
201 a.n() == b.n(),
202 "invalid inputs: a.n() = {} != b.n() = {}",
203 a.n(),
204 b.n()
205 );
206
207 let b_vec: &mut Vec<u64> = &mut b.0;
208 let a_vec: &Vec<u64> = &a.0;
209
210 if NTT {
211 a_vec.iter().enumerate().for_each(|(i, ai)| match MOD {
212 0 => b_vec[idx[i]] = *ai,
213 1 => self
214 .modulus
215 .sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx[i]]),
216 2 => self
217 .modulus
218 .sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx[i]]),
219 _ => {
220 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
221 }
222 });
223 } else {
224 let n: usize = a.n();
225 let mask: usize = n - 1;
226 let q: u64 = self.modulus.q();
227 a_vec.iter().enumerate().for_each(|(i, ai)| {
228 let sign: u64 = (idx[i] >> usize::BITS - 1) as u64;
229 let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
230 match MOD {
231 0 => b_vec[idx[i] & mask] = v,
232 1 => self
233 .modulus
234 .sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[idx[i] & mask]),
235 2 => self
236 .modulus
237 .sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[idx[i] & mask]),
238 _ => {
239 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
240 }
241 }
242 });
243 }
244 }
245}