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