poulpy/math/src/ring/impl_u64/automorphism.rs
math/src/ring/impl_u64/automorphism.rsRaw
1use crate::automorphism::AutomorphismPermutation;
2use crate::modulus::{ScalarOperations, ONCE};
3use crate::modulus::{WordOps, REDUCEMOD};
4use crate::poly::Poly;
5use crate::ring::Ring;
6
7impl Ring<u64> {
8 // b <- auto(a)
9 pub fn a_apply_automorphism_native_into_b<const NTT: bool>(
10 &self,
11 a: &Poly<u64>,
12 gal_el: usize,
13 nth_root: usize,
14 b: &mut Poly<u64>,
15 ) {
16 self.apply_automorphism_native_core::<0, ONCE, NTT>(a, gal_el, nth_root, b)
17 }
18
19 // b <- REDUCEMOD(b + auto(a))
20 pub fn a_apply_automorphism_native_add_b_into_b<const REDUCE: REDUCEMOD, const NTT: bool>(
21 &self,
22 a: &Poly<u64>,
23 gal_el: usize,
24 nth_root: usize,
25 b: &mut Poly<u64>,
26 ) {
27 self.apply_automorphism_native_core::<1, REDUCE, NTT>(a, gal_el, nth_root, b)
28 }
29
30 // b <- REDUCEMOD(b - auto(a))
31 pub fn a_apply_automorphism_native_sub_b_into_b<const REDUCE: REDUCEMOD, const NTT: bool>(
32 &self,
33 a: &Poly<u64>,
34 gal_el: usize,
35 nth_root: usize,
36 b: &mut Poly<u64>,
37 ) {
38 self.apply_automorphism_native_core::<2, REDUCE, NTT>(a, gal_el, nth_root, b)
39 }
40
41 fn apply_automorphism_native_core<const MOD: u8, const REDUCE: REDUCEMOD, const NTT: bool>(
42 &self,
43 a: &Poly<u64>,
44 gal_el: usize,
45 nth_root: usize,
46 b: &mut Poly<u64>,
47 ) {
48 debug_assert!(
49 a.n() == b.n(),
50 "invalid inputs: a.n() = {} != b.n() = {}",
51 a.n(),
52 b.n()
53 );
54
55 assert!(
56 gal_el & 1 == 1,
57 "invalid gal_el={}: not coprime with nth_root={}",
58 gal_el,
59 nth_root
60 );
61
62 assert!(
63 nth_root & (nth_root - 1) == 0,
64 "invalid nth_root={}: not a power-of-two",
65 nth_root
66 );
67
68 let b_vec: &mut Vec<u64> = &mut b.0;
69 let a_vec: &Vec<u64> = &a.0;
70
71 if NTT {
72 let mask: usize = nth_root - 1;
73 let log_nth_root_half: u32 = nth_root.log2() as u32 - 1;
74 a_vec.iter().enumerate().for_each(|(i, ai)| {
75 let i_rev: usize = 2 * i.reverse_bits_msb(log_nth_root_half) + 1;
76 let gal_el_i: usize = (((gal_el * i_rev) & mask) - 1) >> 1;
77 let idx: usize = gal_el_i.reverse_bits_msb(log_nth_root_half);
78 match MOD {
79 0 => b_vec[idx] = *ai,
80 1 => self
81 .modulus
82 .sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx]),
83 2 => self
84 .modulus
85 .sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx]),
86 _ => {
87 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
88 }
89 }
90 });
91 } else {
92 let n: usize = a.n();
93 let mask: usize = n - 1;
94 let log_n: usize = n.log2();
95 let q: u64 = self.modulus.q();
96 a_vec.iter().enumerate().for_each(|(i, ai)| {
97 let gal_el_i: usize = i * gal_el;
98 let sign: u64 = ((gal_el_i >> log_n) & 1) as u64;
99 let i_out: usize = gal_el_i & mask;
100 let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
101 match MOD {
102 0 => b_vec[i_out] = v,
103 1 => self
104 .modulus
105 .sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[i_out]),
106 2 => self
107 .modulus
108 .sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[i_out]),
109 _ => {
110 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
111 }
112 }
113 });
114 }
115 }
116
117 // b <- auto(a)
118 pub fn a_apply_automorphism_from_perm_into_b<const NTT: bool>(
119 &self,
120 a: &Poly<u64>,
121 auto_perm: &AutomorphismPermutation,
122 b: &mut Poly<u64>,
123 ) {
124 self.automorphism_from_perm_core::<0, ONCE, NTT>(a, auto_perm, b)
125 }
126
127 // b <- REDUCEMOD(b + auto(a))
128 pub fn a_apply_automorphism_from_perm_add_b_into_b<const REDUCE: REDUCEMOD, const NTT: bool>(
129 &self,
130 a: &Poly<u64>,
131 auto_perm: &AutomorphismPermutation,
132 b: &mut Poly<u64>,
133 ) {
134 self.automorphism_from_perm_core::<1, REDUCE, NTT>(a, auto_perm, b)
135 }
136
137 // b <- REDUCEMOD(b - auto(a))
138 pub fn a_apply_automorphism_from_perm_sub_b_into_b<const REDUCE: REDUCEMOD, const NTT: bool>(
139 &self,
140 a: &Poly<u64>,
141 auto_perm: &AutomorphismPermutation,
142 b: &mut Poly<u64>,
143 ) {
144 self.automorphism_from_perm_core::<2, REDUCE, NTT>(a, auto_perm, b)
145 }
146
147 // b <- auto(a) if OVERWRITE else b <- REDUCEMOD(b + auto(a))
148 fn automorphism_from_perm_core<const MOD: u8, const REDUCE: REDUCEMOD, const NTT: bool>(
149 &self,
150 a: &Poly<u64>,
151 auto_perm: &AutomorphismPermutation,
152 b: &mut Poly<u64>,
153 ) {
154 debug_assert!(
155 a.n() == b.n(),
156 "invalid inputs: a.n() = {} != b.n() = {}",
157 a.n(),
158 b.n()
159 );
160
161 let b_vec: &mut Vec<u64> = &mut b.0;
162 let a_vec: &Vec<u64> = &a.0;
163
164 let idx: &Vec<usize> = &auto_perm.permutation;
165
166 if NTT {
167 a_vec.iter().enumerate().for_each(|(i, ai)| match MOD {
168 0 => b_vec[idx[i]] = *ai,
169 1 => self
170 .modulus
171 .sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx[i]]),
172 2 => self
173 .modulus
174 .sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx[i]]),
175 _ => {
176 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
177 }
178 });
179 } else {
180 let n: usize = a.n();
181 let mask: usize = n - 1;
182 let q: u64 = self.modulus.q();
183 a_vec.iter().enumerate().for_each(|(i, ai)| {
184 let sign: u64 = (idx[i] >> usize::BITS - 1) as u64;
185 let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
186 match MOD {
187 0 => b_vec[idx[i] & mask] = v,
188 1 => self
189 .modulus
190 .sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[idx[i] & mask]),
191 2 => self
192 .modulus
193 .sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[idx[i] & mask]),
194 _ => {
195 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
196 }
197 }
198 });
199 }
200 }
201}