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    pub 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^|gen| * sign(gen)
 69    pub fn galois_element(&self, gen: i64) -> i64 {
 70        if gen == 0 {
 71            return 1;
 72        }
 73        ((mod_exp_u64(GALOISGENERATOR, gen.abs() as usize) & (self.cyclotomic_order() - 1)) as i64) * gen.signum()
 74    }
 75
 76    // Returns gen^-1
 77    pub fn galois_element_inv(&self, gen: i64) -> i64 {
 78        if gen == 0 {
 79            panic!("cannot invert 0")
 80        }
 81        ((mod_exp_u64(gen.abs() as u64, (self.cyclotomic_order() - 1) as usize) & (self.cyclotomic_order() - 1)) as i64)
 82            * gen.signum()
 83    }
 84
 85    pub fn free(self) {
 86        unsafe { delete_module_info(self.ptr) }
 87        drop(self);
 88    }
 89}
 90
 91fn mod_exp_u64(x: u64, e: usize) -> u64 {
 92    let mut y: u64 = 1;
 93    let mut x_pow: u64 = x;
 94    let mut exp = e;
 95    while exp > 0 {
 96        if exp & 1 == 1 {
 97            y = y.wrapping_mul(x_pow);
 98        }
 99        x_pow = x_pow.wrapping_mul(x_pow);
100        exp >>= 1;
101    }
102    y
103}