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