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