arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/base2k/src/lib.rs
base2k/src/lib.rsRaw
  1pub mod commons;
  2pub mod encoding;
  3#[allow(non_camel_case_types, non_snake_case, non_upper_case_globals, dead_code, improper_ctypes)]
  4// Other modules and exports
  5pub mod ffi;
  6pub mod mat_znx_dft;
  7pub mod module;
  8pub mod sampling;
  9pub mod scalar_znx_dft;
 10pub mod stats;
 11pub mod vec_znx;
 12pub mod vec_znx_big;
 13pub mod vec_znx_dft;
 14pub mod vec_znx_ops;
 15
 16pub use commons::*;
 17pub use encoding::*;
 18pub use mat_znx_dft::*;
 19pub use module::*;
 20pub use sampling::*;
 21pub use scalar_znx_dft::*;
 22#[allow(unused_imports)]
 23pub use stats::*;
 24pub use vec_znx::*;
 25pub use vec_znx_big::*;
 26pub use vec_znx_dft::*;
 27pub use vec_znx_ops::*;
 28
 29pub const GALOISGENERATOR: u64 = 5;
 30pub const DEFAULTALIGN: usize = 64;
 31
 32fn is_aligned_custom<T>(ptr: *const T, align: usize) -> bool {
 33    (ptr as usize) % align == 0
 34}
 35
 36pub fn is_aligned<T>(ptr: *const T) -> bool {
 37    is_aligned_custom(ptr, DEFAULTALIGN)
 38}
 39
 40pub fn assert_alignement<T>(ptr: *const T) {
 41    assert!(
 42        is_aligned(ptr),
 43        "invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]"
 44    )
 45}
 46
 47pub fn cast<T, V>(data: &[T]) -> &[V] {
 48    let ptr: *const V = data.as_ptr() as *const V;
 49    let len: usize = data.len() / std::mem::size_of::<V>();
 50    unsafe { std::slice::from_raw_parts(ptr, len) }
 51}
 52
 53pub fn cast_mut<T, V>(data: &[T]) -> &mut [V] {
 54    let ptr: *mut V = data.as_ptr() as *mut V;
 55    let len: usize = data.len() / std::mem::size_of::<V>();
 56    unsafe { std::slice::from_raw_parts_mut(ptr, len) }
 57}
 58
 59/// Allocates a block of bytes with a custom alignement.
 60/// Alignement must be a power of two and size a multiple of the alignement.
 61/// Allocated memory is initialized to zero.
 62fn alloc_aligned_custom_u8(size: usize, align: usize) -> Vec<u8> {
 63    assert!(
 64        align.is_power_of_two(),
 65        "Alignment must be a power of two but is {}",
 66        align
 67    );
 68    assert_eq!(
 69        (size * std::mem::size_of::<u8>()) % align,
 70        0,
 71        "size={} must be a multiple of align={}",
 72        size,
 73        align
 74    );
 75    unsafe {
 76        let layout: std::alloc::Layout = std::alloc::Layout::from_size_align(size, align).expect("Invalid alignment");
 77        let ptr: *mut u8 = std::alloc::alloc(layout);
 78        if ptr.is_null() {
 79            panic!("Memory allocation failed");
 80        }
 81        assert!(
 82            is_aligned_custom(ptr, align),
 83            "Memory allocation at {:p} is not aligned to {} bytes",
 84            ptr,
 85            align
 86        );
 87        // Init allocated memory to zero
 88        std::ptr::write_bytes(ptr, 0, size);
 89        Vec::from_raw_parts(ptr, size, size)
 90    }
 91}
 92
 93/// Allocates a block of T aligned with [DEFAULTALIGN].
 94/// Size of T * size msut be a multiple of [DEFAULTALIGN].
 95pub fn alloc_aligned_custom<T>(size: usize, align: usize) -> Vec<T> {
 96    assert_eq!(
 97        (size * std::mem::size_of::<T>()) % align,
 98        0,
 99        "size={} must be a multiple of align={}",
100        size,
101        align
102    );
103    let mut vec_u8: Vec<u8> = alloc_aligned_custom_u8(std::mem::size_of::<T>() * size, align);
104    let ptr: *mut T = vec_u8.as_mut_ptr() as *mut T;
105    let len: usize = vec_u8.len() / std::mem::size_of::<T>();
106    let cap: usize = vec_u8.capacity() / std::mem::size_of::<T>();
107    std::mem::forget(vec_u8);
108    unsafe { Vec::from_raw_parts(ptr, len, cap) }
109}
110
111// Allocates an aligned of size equal to the smallest power of two equal or greater to `size` that is
112// at least as bit as DEFAULTALIGN / std::mem::size_of::<T>().
113pub fn alloc_aligned<T>(size: usize) -> Vec<T> {
114    alloc_aligned_custom::<T>(
115        std::cmp::max(
116            size.next_power_of_two(),
117            DEFAULTALIGN / std::mem::size_of::<T>(),
118        ),
119        DEFAULTALIGN,
120    )
121}