1pub mod encoding;
  2#[allow(non_camel_case_types, non_snake_case, non_upper_case_globals, dead_code, improper_ctypes)]
  3// Other modules and exports
  4pub mod ffi;
  5pub mod mat_znx_dft;
  6pub mod mat_znx_dft_ops;
  7pub mod module;
  8pub mod sampling;
  9pub mod scalar_znx;
 10pub mod scalar_znx_dft;
 11pub mod scalar_znx_dft_ops;
 12pub mod stats;
 13pub mod vec_znx;
 14pub mod vec_znx_big;
 15pub mod vec_znx_big_ops;
 16pub mod vec_znx_dft;
 17pub mod vec_znx_dft_ops;
 18pub mod vec_znx_ops;
 19pub mod znx_base;
 20
 21pub use encoding::*;
 22pub use mat_znx_dft::*;
 23pub use mat_znx_dft_ops::*;
 24pub use module::*;
 25pub use sampling::*;
 26pub use scalar_znx::*;
 27pub use scalar_znx_dft::*;
 28pub use scalar_znx_dft_ops::*;
 29pub use stats::*;
 30pub use vec_znx::*;
 31pub use vec_znx_big::*;
 32pub use vec_znx_big_ops::*;
 33pub use vec_znx_dft::*;
 34pub use vec_znx_dft_ops::*;
 35pub use vec_znx_ops::*;
 36pub use znx_base::*;
 37
 38pub const GALOISGENERATOR: u64 = 5;
 39pub const DEFAULTALIGN: usize = 64;
 40
 41fn is_aligned_custom<T>(ptr: *const T, align: usize) -> bool {
 42    (ptr as usize) % align == 0
 43}
 44
 45pub fn is_aligned<T>(ptr: *const T) -> bool {
 46    is_aligned_custom(ptr, DEFAULTALIGN)
 47}
 48
 49pub fn assert_alignement<T>(ptr: *const T) {
 50    assert!(
 51        is_aligned(ptr),
 52        "invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]"
 53    )
 54}
 55
 56pub fn cast<T, V>(data: &[T]) -> &[V] {
 57    let ptr: *const V = data.as_ptr() as *const V;
 58    let len: usize = data.len() / size_of::<V>();
 59    unsafe { std::slice::from_raw_parts(ptr, len) }
 60}
 61
 62pub fn cast_mut<T, V>(data: &[T]) -> &mut [V] {
 63    let ptr: *mut V = data.as_ptr() as *mut V;
 64    let len: usize = data.len() / size_of::<V>();
 65    unsafe { std::slice::from_raw_parts_mut(ptr, len) }
 66}
 67
 68/// Allocates a block of bytes with a custom alignement.
 69/// Alignement must be a power of two and size a multiple of the alignement.
 70/// Allocated memory is initialized to zero.
 71fn alloc_aligned_custom_u8(size: usize, align: usize) -> Vec<u8> {
 72    assert!(
 73        align.is_power_of_two(),
 74        "Alignment must be a power of two but is {}",
 75        align
 76    );
 77    assert_eq!(
 78        (size * size_of::<u8>()) % align,
 79        0,
 80        "size={} must be a multiple of align={}",
 81        size,
 82        align
 83    );
 84    unsafe {
 85        let layout: std::alloc::Layout = std::alloc::Layout::from_size_align(size, align).expect("Invalid alignment");
 86        let ptr: *mut u8 = std::alloc::alloc(layout);
 87        if ptr.is_null() {
 88            panic!("Memory allocation failed");
 89        }
 90        assert!(
 91            is_aligned_custom(ptr, align),
 92            "Memory allocation at {:p} is not aligned to {} bytes",
 93            ptr,
 94            align
 95        );
 96        // Init allocated memory to zero
 97        std::ptr::write_bytes(ptr, 0, size);
 98        Vec::from_raw_parts(ptr, size, size)
 99    }
100}
101
102/// Allocates a block of T aligned with [DEFAULTALIGN].
103/// Size of T * size msut be a multiple of [DEFAULTALIGN].
104pub fn alloc_aligned_custom<T>(size: usize, align: usize) -> Vec<T> {
105    assert_eq!(
106        (size * size_of::<T>()) % align,
107        0,
108        "size={} must be a multiple of align={}",
109        size,
110        align
111    );
112    let mut vec_u8: Vec<u8> = alloc_aligned_custom_u8(size_of::<T>() * size, align);
113    let ptr: *mut T = vec_u8.as_mut_ptr() as *mut T;
114    let len: usize = vec_u8.len() / size_of::<T>();
115    let cap: usize = vec_u8.capacity() / size_of::<T>();
116    std::mem::forget(vec_u8);
117    unsafe { Vec::from_raw_parts(ptr, len, cap) }
118}
119
120/// Allocates an aligned vector of size equal to the smallest multiple
121/// of [DEFAULTALIGN]/size_of::<T>() that is equal or greater to `size`.
122pub fn alloc_aligned<T>(size: usize) -> Vec<T> {
123    alloc_aligned_custom::<T>(
124        size + (size % (DEFAULTALIGN / size_of::<T>())),
125        DEFAULTALIGN,
126    )
127}
128
129// Scratch implementation below
130
131pub struct ScratchOwned(Vec<u8>);
132
133impl ScratchOwned {
134    pub fn new(byte_count: usize) -> Self {
135        let data: Vec<u8> = alloc_aligned(byte_count);
136        Self(data)
137    }
138
139    pub fn borrow(&mut self) -> &mut Scratch {
140        Scratch::new(&mut self.0)
141    }
142}
143
144pub struct Scratch {
145    data: [u8],
146}
147
148impl Scratch {
149    fn new(data: &mut [u8]) -> &mut Self {
150        unsafe { &mut *(data as *mut [u8] as *mut Self) }
151    }
152
153    fn take_slice_aligned(data: &mut [u8], take_len: usize) -> (&mut [u8], &mut [u8]) {
154        let ptr = data.as_mut_ptr();
155        let self_len = data.len();
156
157        let aligned_offset = ptr.align_offset(DEFAULTALIGN);
158        let aligned_len = self_len.saturating_sub(aligned_offset);
159
160        if let Some(rem_len) = aligned_len.checked_sub(take_len) {
161            unsafe {
162                let rem_ptr = ptr.add(aligned_offset).add(take_len);
163                let rem_slice = &mut *std::ptr::slice_from_raw_parts_mut(rem_ptr, rem_len);
164
165                let take_slice = &mut *std::ptr::slice_from_raw_parts_mut(ptr.add(aligned_offset), take_len);
166
167                return (take_slice, rem_slice);
168            }
169        } else {
170            panic!(
171                "Attempted to take {} from scratch with {} aligned bytes left",
172                take_len,
173                aligned_len,
174                // type_name::<T>(),
175                // aligned_len
176            );
177        }
178    }
179
180    pub fn tmp_scalar_slice<T>(&mut self, len: usize) -> (&mut [T], &mut Self) {
181        let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, len * std::mem::size_of::<T>());
182
183        unsafe {
184            (
185                &mut *(std::ptr::slice_from_raw_parts_mut(take_slice.as_mut_ptr() as *mut T, len)),
186                Self::new(rem_slice),
187            )
188        }
189    }
190
191    pub fn tmp_vec_znx_dft<B: Backend>(
192        &mut self,
193        module: &Module<B>,
194        cols: usize,
195        size: usize,
196    ) -> (VecZnxDft<&mut [u8], B>, &mut Self) {
197        let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_vec_znx_dft(module, cols, size));
198
199        (
200            VecZnxDft::from_data(take_slice, module.n(), cols, size),
201            Self::new(rem_slice),
202        )
203    }
204
205    pub fn tmp_vec_znx_big<B: Backend>(
206        &mut self,
207        module: &Module<B>,
208        cols: usize,
209        size: usize,
210    ) -> (VecZnxBig<&mut [u8], B>, &mut Self) {
211        let (take_slice, rem_slice) = Self::take_slice_aligned(&mut self.data, bytes_of_vec_znx_big(module, cols, size));
212
213        (
214            VecZnxBig::from_data(take_slice, module.n(), cols, size),
215            Self::new(rem_slice),
216        )
217    }
218}