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
37
38pub fn is_aligned_custom<T>(ptr: *const T, align: usize) -> bool {
39 (ptr as usize) % align == 0
40}
41
42pub fn is_aligned<T>(ptr: *const T) -> bool {
43 is_aligned_custom(ptr, DEFAULTALIGN)
44}
45
46pub fn assert_alignement<T>(ptr: *const T) {
47 assert!(
48 is_aligned(ptr),
49 "invalid alignement: ensure passed bytes have been allocated with [alloc_aligned_u8] or [alloc_aligned]"
50 )
51}
52
53pub fn cast<T, V>(data: &[T]) -> &[V] {
54 let ptr: *const V = data.as_ptr() as *const V;
55 let len: usize = data.len() / std::mem::size_of::<V>();
56 unsafe { std::slice::from_raw_parts(ptr, len) }
57}
58
59pub fn cast_mut<T, V>(data: &[T]) -> &mut [V] {
60 let ptr: *mut V = data.as_ptr() as *mut V;
61 let len: usize = data.len() / std::mem::size_of::<V>();
62 unsafe { std::slice::from_raw_parts_mut(ptr, len) }
63}
64
65use std::alloc::{Layout, alloc};
66use std::ptr;
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.
71pub fn 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 * std::mem::size_of::<u8>()) % align,
79 0,
80 "size={} must be a multiple of align={}",
81 size,
82 align
83 );
84 unsafe {
85 let layout: Layout = Layout::from_size_align(size, align).expect("Invalid alignment");
86 let ptr: *mut u8 = 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 ptr::write_bytes(ptr, 0, size);
98 Vec::from_raw_parts(ptr, size, size)
99 }
100}
101
102/// Allocates a block of bytes aligned with [DEFAULTALIGN].
103/// Size must be amultiple of [DEFAULTALIGN].
104/// /// Allocated memory is initialized to zero.
105pub fn alloc_aligned_u8(size: usize) -> Vec<u8> {
106 alloc_aligned_custom_u8(size, DEFAULTALIGN)
107}
108
109/// Allocates a block of T aligned with [DEFAULTALIGN].
110/// Size of T * size msut be a multiple of [DEFAULTALIGN].
111pub fn alloc_aligned_custom<T>(size: usize, align: usize) -> Vec<T> {
112 assert_eq!(
113 (size * std::mem::size_of::<T>()) % align,
114 0,
115 "size={} must be a multiple of align={}",
116 size,
117 align
118 );
119 let mut vec_u8: Vec<u8> = alloc_aligned_custom_u8(std::mem::size_of::<T>() * size, align);
120 let ptr: *mut T = vec_u8.as_mut_ptr() as *mut T;
121 let len: usize = vec_u8.len() / std::mem::size_of::<T>();
122 let cap: usize = vec_u8.capacity() / std::mem::size_of::<T>();
123 std::mem::forget(vec_u8);
124 unsafe { Vec::from_raw_parts(ptr, len, cap) }
125}
126
127pub fn alloc_aligned<T>(size: usize) -> Vec<T> {
128 alloc_aligned_custom::<T>(size, DEFAULTALIGN)
129}