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