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