1use crate::modulus::{REDUCEMOD, WordOps};
  2use crate::poly::Poly;
  3use crate::ring::Ring;
  4use crate::modulus::{ONCE, ScalarOperations};
  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
 49    // b <- auto(a)
 50    pub fn a_apply_automorphism_into_b<const NTT: bool>(
 51        &self,
 52        a: &Poly<u64>,
 53        gal_el: usize,
 54        nth_root: usize,
 55        b: &mut Poly<u64>,
 56    ){
 57        self.apply_automorphism_core::<0, ONCE, NTT>(a, gal_el, nth_root, b)
 58    }
 59
 60    // b <- REDUCEMOD(b + auto(a))
 61    pub fn a_apply_automorphism_add_b_into_b<const REDUCE:REDUCEMOD, const NTT: bool>(
 62        &self,
 63        a: &Poly<u64>,
 64        gal_el: usize,
 65        nth_root: usize,
 66        b: &mut Poly<u64>,
 67    ){
 68        self.apply_automorphism_core::<1, REDUCE, NTT>(a, gal_el, nth_root, b)
 69    }
 70
 71    // b <- REDUCEMOD(b - auto(a))
 72    pub fn a_apply_automorphism_sub_b_into_b<const REDUCE:REDUCEMOD, const NTT:bool>(
 73    &self,
 74    a: &Poly<u64>,
 75    gal_el: usize,
 76    nth_root: usize,
 77    b: &mut Poly<u64>,
 78    ){
 79        self.apply_automorphism_core::<2, REDUCE, NTT>(a, gal_el, nth_root, b)
 80    }
 81
 82    fn apply_automorphism_core<const MOD:u8, const REDUCE:REDUCEMOD, const NTT: bool>(
 83        &self,
 84        a: &Poly<u64>,
 85        gal_el: usize,
 86        nth_root: usize,
 87        b: &mut Poly<u64>,
 88    ) {
 89        debug_assert!(
 90            a.n() == b.n(),
 91            "invalid inputs: a.n() = {} != b.n() = {}",
 92            a.n(),
 93            b.n()
 94        );
 95
 96        assert!(
 97            gal_el & 1 == 1,
 98            "invalid gal_el={}: not coprime with nth_root={}",
 99            gal_el,
100            nth_root
101        );
102
103        assert!(
104            nth_root & (nth_root - 1) == 0,
105            "invalid nth_root={}: not a power-of-two",
106            nth_root
107        );
108
109        let b_vec: &mut Vec<u64> = &mut b.0;
110        let a_vec: &Vec<u64> = &a.0;
111
112        if NTT {
113            let mask: usize = nth_root - 1;
114            let log_nth_root_half: u32 = nth_root.log2() as u32 - 1;
115            a_vec.iter().enumerate().for_each(|(i, ai)| {
116                let i_rev: usize = 2 * i.reverse_bits_msb(log_nth_root_half) + 1;
117                let gal_el_i: usize = (((gal_el * i_rev) & mask) - 1) >> 1;
118                let idx: usize = gal_el_i.reverse_bits_msb(log_nth_root_half);
119                match MOD{
120                    0 =>{b_vec[idx] = *ai}
121                    1=>{self.modulus.sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx])}
122                    2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx])}
123                    _=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
124                }
125            });
126        } else {
127            let n: usize = a.n();
128            let mask: usize = n - 1;
129            let log_n: usize = n.log2();
130            let q: u64 = self.modulus.q();
131            a_vec.iter().enumerate().for_each(|(i, ai)| {
132                let gal_el_i: usize = i * gal_el;
133                let sign: u64 = ((gal_el_i >> log_n) & 1) as u64;
134                let i_out: usize = gal_el_i & mask;
135                let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
136                match MOD{
137                    0 =>{b_vec[i_out] = v}
138                    1=>{self.modulus.sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[i_out])}
139                    2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[i_out])}
140                    _=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
141                }
142            });
143        }
144    }
145
146    // b <- auto(a)
147    pub fn a_apply_automorphism_from_index_into_b<const NTT: bool>(
148        &self,
149        a: &Poly<u64>,
150        idx: &[usize],
151        b: &mut Poly<u64>,
152    ){
153        self.automorphism_from_index_core::<0, ONCE, NTT>(a, idx, b)
154    }
155
156    // b <- REDUCEMOD(b + auto(a))
157    pub fn a_apply_automorphism_from_index_add_b_into_b<const REDUCE:REDUCEMOD, const NTT: bool>(
158        &self,
159        a: &Poly<u64>,
160        idx: &[usize],
161        b: &mut Poly<u64>,
162    ){
163        self.automorphism_from_index_core::<1, REDUCE, NTT>(a, idx, b)
164    }
165
166    // b <- REDUCEMOD(b - auto(a))
167    pub fn a_apply_automorphism_from_index_sub_b_into_b<const REDUCE:REDUCEMOD, const NTT:bool>(
168    &self,
169    a: &Poly<u64>,
170    idx: &[usize],
171    b: &mut Poly<u64>,
172    ){
173        self.automorphism_from_index_core::<2, REDUCE, NTT>(a, idx, b)
174    }
175
176    // b <- auto(a) if OVERWRITE else b <- REDUCEMOD(b + auto(a))
177    fn automorphism_from_index_core<const MOD:u8, const REDUCE:REDUCEMOD, const NTT: bool>(
178        &self,
179        a: &Poly<u64>,
180        idx: &[usize],
181        b: &mut Poly<u64>,
182    ) {
183        debug_assert!(
184            a.n() == b.n(),
185            "invalid inputs: a.n() = {} != b.n() = {}",
186            a.n(),
187            b.n()
188        );
189
190        let b_vec: &mut Vec<u64> = &mut b.0;
191        let a_vec: &Vec<u64> = &a.0;
192
193        if NTT {
194            a_vec.iter().enumerate().for_each(|(i, ai)| {
195                match MOD{
196                    0 =>{ b_vec[idx[i]] = *ai}
197                    1 =>{self.modulus.sa_add_sb_into_sb::<REDUCE>(ai, &mut b_vec[idx[i]])}
198                    2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(ai, &mut b_vec[idx[i]])}
199                    _=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
200                }
201            });
202        } else {
203            let n: usize = a.n();
204            let mask: usize = n - 1;
205            let q: u64 = self.modulus.q();
206            a_vec.iter().enumerate().for_each(|(i, ai)| {
207                let sign: u64 = (idx[i] >> usize::BITS - 1) as u64;
208                let v: u64 = ai * (sign ^ 1) | (q - ai) * sign;
209                match MOD{
210                    0 =>{b_vec[idx[i] & mask] = v}
211                    1 =>{self.modulus.sa_add_sb_into_sb::<REDUCE>(&v, &mut b_vec[idx[i] & mask])}
212                    2=>{self.modulus.sa_sub_sb_into_sa::<1, REDUCE>(&v, &mut b_vec[idx[i] & mask])}
213                    _=>{panic!("invalid const MOD should be 0, 1, or 2 but is {}", MOD)}
214                }
215            });
216        }
217    }
218}