1use std::{
  2    fmt::{Debug, Display},
  3    marker::PhantomData,
  4    ptr::NonNull,
  5};
  6
  7use rand_distr::num_traits::Zero;
  8
  9use crate::{
 10    GALOISGENERATOR,
 11    api::{ModuleLogN, ModuleN},
 12};
 13
 14#[allow(clippy::missing_safety_doc)]
 15pub trait Backend: Sized {
 16    type ScalarBig: Copy + Zero + Display + Debug;
 17    type ScalarPrep: Copy + Zero + Display + Debug;
 18    type Handle: 'static;
 19    fn layout_prep_word_count() -> usize;
 20    fn layout_big_word_count() -> usize;
 21    unsafe fn destroy(handle: NonNull<Self::Handle>);
 22}
 23
 24pub struct Module<B: Backend> {
 25    ptr: NonNull<B::Handle>,
 26    n: u64,
 27    _marker: PhantomData<B>,
 28}
 29
 30unsafe impl<B: Backend> Sync for Module<B> {}
 31unsafe impl<B: Backend> Send for Module<B> {}
 32
 33impl<B: Backend> Module<B> {
 34    #[allow(clippy::missing_safety_doc)]
 35    #[inline]
 36    pub fn new_marker(n: u64) -> Self {
 37        Self {
 38            ptr: NonNull::dangling(),
 39            n,
 40            _marker: PhantomData,
 41        }
 42    }
 43
 44    #[allow(clippy::missing_safety_doc)]
 45    #[inline]
 46    pub unsafe fn from_nonnull(ptr: NonNull<B::Handle>, n: u64) -> Self {
 47        Self {
 48            ptr,
 49            n,
 50            _marker: PhantomData,
 51        }
 52    }
 53
 54    /// Construct from a raw pointer managed elsewhere.
 55    /// SAFETY: `ptr` must be non-null and remain valid for the lifetime of this Module.
 56    #[inline]
 57    #[allow(clippy::missing_safety_doc)]
 58    pub unsafe fn from_raw_parts(ptr: *mut B::Handle, n: u64) -> Self {
 59        Self {
 60            ptr: NonNull::new(ptr).expect("null module ptr"),
 61            n,
 62            _marker: PhantomData,
 63        }
 64    }
 65
 66    #[allow(clippy::missing_safety_doc)]
 67    #[inline]
 68    pub unsafe fn ptr(&self) -> *mut <B as Backend>::Handle {
 69        self.ptr.as_ptr()
 70    }
 71
 72    #[inline]
 73    pub fn n(&self) -> usize {
 74        self.n as usize
 75    }
 76    #[inline]
 77    pub fn as_mut_ptr(&self) -> *mut B::Handle {
 78        self.ptr.as_ptr()
 79    }
 80
 81    #[inline]
 82    pub fn log_n(&self) -> usize {
 83        (usize::BITS - (self.n() - 1).leading_zeros()) as _
 84    }
 85}
 86
 87pub trait CyclotomicOrder
 88where
 89    Self: ModuleN,
 90{
 91    fn cyclotomic_order(&self) -> i64 {
 92        (self.n() << 1) as _
 93    }
 94}
 95
 96impl<BE: Backend> ModuleLogN for Module<BE> where Self: ModuleN {}
 97
 98impl<BE: Backend> CyclotomicOrder for Module<BE> where Self: ModuleN {}
 99
100pub trait GaloisElement
101where
102    Self: CyclotomicOrder,
103{
104    // Returns GALOISGENERATOR^|generator| * sign(generator)
105    fn galois_element(&self, generator: i64) -> i64 {
106        if generator == 0 {
107            return 1;
108        }
109
110        let g_exp: u64 = mod_exp_u64(GALOISGENERATOR, generator.unsigned_abs() as usize) & (self.cyclotomic_order() - 1) as u64;
111        g_exp as i64 * generator.signum()
112    }
113
114    // Returns gen^-1
115    fn galois_element_inv(&self, gal_el: i64) -> i64 {
116        if gal_el == 0 {
117            panic!("cannot invert 0")
118        }
119
120        let g_exp: u64 = mod_exp_u64(
121            gal_el.unsigned_abs(),
122            (self.cyclotomic_order() - 1) as usize,
123        ) & (self.cyclotomic_order() - 1) as u64;
124        g_exp as i64 * gal_el.signum()
125    }
126}
127
128impl<BE: Backend> GaloisElement for Module<BE> where Self: CyclotomicOrder {}
129
130impl<B: Backend> Drop for Module<B> {
131    fn drop(&mut self) {
132        unsafe { B::destroy(self.ptr) }
133    }
134}
135
136pub fn mod_exp_u64(x: u64, e: usize) -> u64 {
137    let mut y: u64 = 1;
138    let mut x_pow: u64 = x;
139    let mut exp = e;
140    while exp > 0 {
141        if exp & 1 == 1 {
142            y = y.wrapping_mul(x_pow);
143        }
144        x_pow = x_pow.wrapping_mul(x_pow);
145        exp >>= 1;
146    }
147    y
148}