poulpy/math/src/ring/impl_u64/automorphism.rs
math/src/ring/impl_u64/automorphism.rsRaw
1use crate::automorphism::AutoPerm;
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: &AutoPerm,
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: &AutoPerm,
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: &AutoPerm,
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: &AutoPerm,
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 assert!(
162 NTT == auto_perm.ntt,
163 "missmatch between AutoPerm NTT flag={} and method NTT flag={}",
164 auto_perm.ntt,
165 NTT
166 );
167
168 let b_vec: &mut Vec<u64> = &mut b.0;
169 let a_vec: &Vec<u64> = &a.0;
170
171 let idx: &Vec<usize> = &auto_perm.permutation;
172
173 if NTT {
174 a_vec.iter().enumerate().for_each(|(i, ai)| match MOD {
175 0 => b_vec[idx[i]] = *ai,
176 1 => self
177 .modulus
178 .sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx[i]]),
179 2 => self
180 .modulus
181 .sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx[i]]),
182 _ => {
183 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
184 }
185 });
186 } else {
187 let n: usize = a.n();
188 let mask: usize = n - 1;
189 let q: u64 = self.modulus.q();
190 a_vec.iter().enumerate().for_each(|(i, ai)| {
191 let sign: u64 = (idx[i] >> usize::BITS - 1) as u64;
192 let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
193 match MOD {
194 0 => b_vec[idx[i] & mask] = v,
195 1 => self
196 .modulus
197 .sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[idx[i] & mask]),
198 2 => self
199 .modulus
200 .sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[idx[i] & mask]),
201 _ => {
202 panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)
203 }
204 }
205 });
206 }
207 }
208}