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