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