1use crate::modulus::WordOps;
  2use crate::poly::Poly;
  3use crate::ring::Ring;
  4
  5/// Returns a lookup table for the automorphism X^{i} -> X^{i * k mod nth_root}.
  6/// Method will panic if n or nth_root are not power-of-two.
  7/// Method will panic if gal_el is not coprime with nth_root.
  8pub fn automorphism_index<const NTT: bool>(n: usize, nth_root: usize, gal_el: usize) -> Vec<usize> {
  9    assert!(n & (n - 1) != 0, "invalid n={}: not a power-of-two", n);
 10    assert!(
 11        nth_root & (nth_root - 1) == 0,
 12        "invalid nth_root={}: not a power-of-two",
 13        n
 14    );
 15    assert!(
 16        gal_el & 1 == 1,
 17        "invalid gal_el={}: not coprime with nth_root={}",
 18        gal_el,
 19        nth_root
 20    );
 21
 22    let mut index: Vec<usize> = Vec::with_capacity(n);
 23
 24    if NTT {
 25        let mask = nth_root - 1;
 26        let log_nth_root_half: u32 = nth_root.log2() as u32 - 1;
 27        for i in 0..n {
 28            let i_rev: usize = 2 * i.reverse_bits_msb(log_nth_root_half) + 1;
 29            let gal_el_i: usize = ((gal_el * i_rev) & mask) >> 1;
 30            index.push(gal_el_i.reverse_bits_msb(log_nth_root_half));
 31        }
 32    } else {
 33        let log_n: usize = n.log2();
 34        let mask: usize = (n - 1) as usize;
 35        for i in 0..n {
 36            let gal_el_i: usize = i as usize * gal_el;
 37            let sign: usize = (gal_el_i >> log_n) & 1;
 38            let i_out: usize = (gal_el_i & mask) | (sign << (usize::BITS - 1));
 39            index.push(i_out)
 40        }
 41    }
 42
 43    index
 44}
 45
 46impl Ring<u64> {
 47    pub fn automorphism<const NTT: bool>(
 48        &self,
 49        a: &Poly<u64>,
 50        gal_el: usize,
 51        nth_root: usize,
 52        b: &mut Poly<u64>,
 53    ) {
 54        debug_assert!(
 55            a.n() == b.n(),
 56            "invalid inputs: a.n() = {} != b.n() = {}",
 57            a.n(),
 58            b.n()
 59        );
 60
 61        assert!(
 62            gal_el & 1 == 1,
 63            "invalid gal_el={}: not coprime with nth_root={}",
 64            gal_el,
 65            nth_root
 66        );
 67
 68        assert!(
 69            nth_root & (nth_root - 1) == 0,
 70            "invalid nth_root={}: not a power-of-two",
 71            nth_root
 72        );
 73
 74        let b_vec: &mut Vec<u64> = &mut b.0;
 75        let a_vec: &Vec<u64> = &a.0;
 76
 77        if NTT {
 78            let mask: usize = nth_root - 1;
 79            let log_nth_root_half: u32 = nth_root.log2() as u32 - 1;
 80            a_vec.iter().enumerate().for_each(|(i, ai)| {
 81                let i_rev: usize = 2 * i.reverse_bits_msb(log_nth_root_half) + 1;
 82                let gal_el_i: usize = (((gal_el * i_rev) & mask) - 1) >> 1;
 83                let idx: usize = gal_el_i.reverse_bits_msb(log_nth_root_half);
 84                b_vec[idx] = *ai;
 85            });
 86        } else {
 87            let n: usize = a.n();
 88            let mask: usize = n - 1;
 89            let log_n: usize = n.log2();
 90            let q: u64 = self.modulus.q();
 91            a_vec.iter().enumerate().for_each(|(i, ai)| {
 92                let gal_el_i: usize = i * gal_el;
 93                let sign: u64 = ((gal_el_i >> log_n) & 1) as u64;
 94                let i_out: usize = gal_el_i & mask;
 95                b_vec[i_out] = ai * (sign ^ 1) | (q - ai) * sign
 96            });
 97        }
 98    }
 99
100    pub fn automorphism_from_index<const NTT: bool>(
101        &self,
102        a: &Poly<u64>,
103        idx: &[usize],
104        b: &mut Poly<u64>,
105    ) {
106        debug_assert!(
107            a.n() == b.n(),
108            "invalid inputs: a.n() = {} != b.n() = {}",
109            a.n(),
110            b.n()
111        );
112
113        let b_vec: &mut Vec<u64> = &mut b.0;
114        let a_vec: &Vec<u64> = &a.0;
115
116        if NTT {
117            a_vec.iter().enumerate().for_each(|(i, ai)| {
118                b_vec[idx[i]] = *ai;
119            });
120        } else {
121            let n: usize = a.n();
122            let mask: usize = n - 1;
123            let q: u64 = self.modulus.q();
124            a_vec.iter().enumerate().for_each(|(i, ai)| {
125                let sign: u64 = (idx[i] >> usize::BITS - 1) as u64;
126                b_vec[idx[i] & mask] = ai * (sign ^ 1) | (q - ai) * sign;
127            });
128        }
129    }
130}