1use std::{marker::PhantomData, ptr::NonNull};
 2
 3use crate::GALOISGENERATOR;
 4
 5pub trait Backend: Sized {
 6    type Handle: 'static;
 7    unsafe fn destroy(handle: NonNull<Self::Handle>);
 8}
 9
10pub struct Module<B: Backend> {
11    ptr: NonNull<B::Handle>,
12    n: u64,
13    _marker: PhantomData<B>,
14}
15
16impl<B: Backend> Module<B> {
17    /// Construct from a raw pointer managed elsewhere.
18    /// SAFETY: `ptr` must be non-null and remain valid for the lifetime of this Module.
19    #[inline]
20    pub unsafe fn from_raw_parts(ptr: *mut B::Handle, n: u64) -> Self {
21        Self {
22            ptr: NonNull::new(ptr).expect("null module ptr"),
23            n,
24            _marker: PhantomData,
25        }
26    }
27
28    #[inline]
29    pub unsafe fn ptr(&self) -> *mut <B as Backend>::Handle {
30        self.ptr.as_ptr()
31    }
32
33    #[inline]
34    pub fn n(&self) -> usize {
35        self.n as usize
36    }
37    #[inline]
38    pub fn as_mut_ptr(&self) -> *mut B::Handle {
39        self.ptr.as_ptr()
40    }
41
42    #[inline]
43    pub fn log_n(&self) -> usize {
44        (usize::BITS - (self.n() - 1).leading_zeros()) as _
45    }
46
47    #[inline]
48    pub fn cyclotomic_order(&self) -> u64 {
49        (self.n() << 1) as _
50    }
51
52    // Returns GALOISGENERATOR^|generator| * sign(generator)
53    #[inline]
54    pub fn galois_element(&self, generator: i64) -> i64 {
55        if generator == 0 {
56            return 1;
57        }
58        ((mod_exp_u64(GALOISGENERATOR, generator.abs() as usize) & (self.cyclotomic_order() - 1)) as i64) * generator.signum()
59    }
60
61    // Returns gen^-1
62    #[inline]
63    pub fn galois_element_inv(&self, gal_el: i64) -> i64 {
64        if gal_el == 0 {
65            panic!("cannot invert 0")
66        }
67        ((mod_exp_u64(gal_el.abs() as u64, (self.cyclotomic_order() - 1) as usize) & (self.cyclotomic_order() - 1)) as i64)
68            * gal_el.signum()
69    }
70}
71
72impl<B: Backend> Drop for Module<B> {
73    fn drop(&mut self) {
74        unsafe { B::destroy(self.ptr) }
75    }
76}
77
78pub fn mod_exp_u64(x: u64, e: usize) -> u64 {
79    let mut y: u64 = 1;
80    let mut x_pow: u64 = x;
81    let mut exp = e;
82    while exp > 0 {
83        if exp & 1 == 1 {
84            y = y.wrapping_mul(x_pow);
85        }
86        x_pow = x_pow.wrapping_mul(x_pow);
87        exp >>= 1;
88    }
89    y
90}