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