1use std::{fmt::Display, marker::PhantomData, ptr::NonNull};
  2
  3use rand_distr::num_traits::Zero;
  4
  5use crate::{api::{ModuleLogN, ModuleN}, 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
 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> ModuleLogN for Module<BE> where Self: ModuleN{}
 90
 91impl<BE: Backend> CyclotomicOrder for Module<BE> where Self: ModuleN {}
 92
 93pub trait GaloisElement
 94where
 95    Self: CyclotomicOrder,
 96{
 97    // Returns GALOISGENERATOR^|generator| * sign(generator)
 98    fn galois_element(&self, generator: i64) -> i64 {
 99        if generator == 0 {
100            return 1;
101        }
102
103        let g_exp: u64 = mod_exp_u64(GALOISGENERATOR, generator.unsigned_abs() as usize) & (self.cyclotomic_order() - 1) as u64;
104        g_exp as i64 * generator.signum()
105    }
106
107    // Returns gen^-1
108    fn galois_element_inv(&self, gal_el: i64) -> i64 {
109        if gal_el == 0 {
110            panic!("cannot invert 0")
111        }
112
113        let g_exp: u64 =
114            mod_exp_u64(GALOISGENERATOR, (self.cyclotomic_order() - 1) as usize) & (self.cyclotomic_order() - 1) as u64;
115        g_exp as i64 * gal_el.signum()
116    }
117}
118
119impl<BE: Backend> GaloisElement for Module<BE> where Self: CyclotomicOrder {}
120
121impl<B: Backend> Drop for Module<B> {
122    fn drop(&mut self) {
123        unsafe { B::destroy(self.ptr) }
124    }
125}
126
127pub fn mod_exp_u64(x: u64, e: usize) -> u64 {
128    let mut y: u64 = 1;
129    let mut x_pow: u64 = x;
130    let mut exp = e;
131    while exp > 0 {
132        if exp & 1 == 1 {
133            y = y.wrapping_mul(x_pow);
134        }
135        x_pow = x_pow.wrapping_mul(x_pow);
136        exp >>= 1;
137    }
138    y
139}