1use crate::GALOISGENERATOR;
  2use crate::ffi::module::{MODULE, delete_module_info, module_info_t, new_module_info};
  3use std::marker::PhantomData;
  4
  5#[derive(Copy, Clone)]
  6#[repr(u8)]
  7pub enum BACKEND {
  8    FFT64,
  9    NTT120,
 10}
 11
 12pub trait Backend {
 13    const KIND: BACKEND;
 14    fn module_type() -> u32;
 15}
 16
 17pub struct FFT64;
 18pub struct NTT120;
 19
 20impl Backend for FFT64 {
 21    const KIND: BACKEND = BACKEND::FFT64;
 22    fn module_type() -> u32 {
 23        0
 24    }
 25}
 26
 27impl Backend for NTT120 {
 28    const KIND: BACKEND = BACKEND::NTT120;
 29    fn module_type() -> u32 {
 30        1
 31    }
 32}
 33
 34pub struct Module<B: Backend> {
 35    pub ptr: *mut MODULE,
 36    n: usize,
 37    _marker: PhantomData<B>,
 38}
 39
 40impl<B: Backend> Module<B> {
 41    // Instantiates a new module.
 42    pub fn new(n: usize) -> Self {
 43        unsafe {
 44            let m: *mut module_info_t = new_module_info(n as u64, B::module_type());
 45            if m.is_null() {
 46                panic!("Failed to create module.");
 47            }
 48            Self {
 49                ptr: m,
 50                n: n,
 51                _marker: PhantomData,
 52            }
 53        }
 54    }
 55
 56    pub fn n(&self) -> usize {
 57        self.n
 58    }
 59
 60    pub fn log_n(&self) -> usize {
 61        (usize::BITS - (self.n() - 1).leading_zeros()) as _
 62    }
 63
 64    pub fn cyclotomic_order(&self) -> u64 {
 65        (self.n() << 1) as _
 66    }
 67
 68    // Returns GALOISGENERATOR^|generator| * sign(generator)
 69    pub fn galois_element(&self, generator: i64) -> i64 {
 70        if generator == 0 {
 71            return 1;
 72        }
 73        ((mod_exp_u64(GALOISGENERATOR, generator.abs() as usize) & (self.cyclotomic_order() - 1)) as i64) * generator.signum()
 74    }
 75
 76    // Returns gen^-1
 77    pub fn galois_element_inv(&self, gal_el: i64) -> i64 {
 78        if gal_el == 0 {
 79            panic!("cannot invert 0")
 80        }
 81        ((mod_exp_u64(gal_el.abs() as u64, (self.cyclotomic_order() - 1) as usize) & (self.cyclotomic_order() - 1)) as i64)
 82            * gal_el.signum()
 83    }
 84}
 85
 86impl<B: Backend> Drop for Module<B> {
 87    fn drop(&mut self) {
 88        unsafe { delete_module_info(self.ptr) }
 89    }
 90}
 91
 92fn mod_exp_u64(x: u64, e: usize) -> u64 {
 93    let mut y: u64 = 1;
 94    let mut x_pow: u64 = x;
 95    let mut exp = e;
 96    while exp > 0 {
 97        if exp & 1 == 1 {
 98            y = y.wrapping_mul(x_pow);
 99        }
100        x_pow = x_pow.wrapping_mul(x_pow);
101        exp >>= 1;
102    }
103    y
104}