croaring/rust_alloc/
global_alloc.rs

1use super::{layout, AlignedLayout};
2
3unsafe extern "C" fn malloc(size: usize) -> *mut core::ffi::c_void {
4    let Some(layout) = layout(size) else {
5        return core::ptr::null_mut();
6    };
7    let ptr = alloc::alloc::alloc(layout);
8    if ptr.is_null() {
9        return ptr.cast();
10    }
11    let size_ptr = ptr.cast::<usize>();
12    size_ptr.write(size);
13    size_ptr.add(1).cast()
14}
15
16unsafe extern "C" fn calloc(nmemb: usize, size: usize) -> *mut core::ffi::c_void {
17    let Some(total_size) = nmemb.checked_mul(size) else {
18        return core::ptr::null_mut();
19    };
20    let Some(layout) = layout(total_size) else {
21        return core::ptr::null_mut();
22    };
23    let ptr = alloc::alloc::alloc_zeroed(layout);
24    if ptr.is_null() {
25        return core::ptr::null_mut();
26    }
27    let size_ptr = ptr.cast::<usize>();
28    size_ptr.write(total_size);
29    size_ptr.add(1).cast()
30}
31
32unsafe extern "C" fn realloc(ptr: *mut core::ffi::c_void, size: usize) -> *mut core::ffi::c_void {
33    if ptr.is_null() {
34        return malloc(size);
35    }
36    let size_ptr = ptr.cast::<usize>().sub(1);
37    let old_size = size_ptr.read();
38    let old_layout = layout(old_size).unwrap();
39    if size == 0 {
40        alloc::alloc::dealloc(size_ptr.cast::<u8>(), old_layout);
41        return core::ptr::null_mut();
42    }
43    let new_ptr = alloc::alloc::realloc(size_ptr.cast(), old_layout, size + size_of::<usize>());
44    if new_ptr.is_null() {
45        return core::ptr::null_mut();
46    }
47    let new_size_ptr = new_ptr.cast::<usize>();
48    new_size_ptr.write(size);
49    new_size_ptr.add(1).cast()
50}
51
52unsafe extern "C" fn free(ptr: *mut core::ffi::c_void) {
53    if ptr.is_null() {
54        return;
55    }
56    let size_ptr = ptr.cast::<usize>().sub(1);
57    let size = size_ptr.read();
58    // If the size would overflow, it would have failed to be allocated in the first place.
59    let layout = layout(size).unwrap();
60    alloc::alloc::dealloc(size_ptr.cast(), layout);
61}
62
63unsafe extern "C" fn aligned_malloc(align: usize, size: usize) -> *mut core::ffi::c_void {
64    let Some(layout) = AlignedLayout::new(size, align) else {
65        return core::ptr::null_mut();
66    };
67    let allocated_ptr = alloc::alloc::alloc(layout.0);
68    if allocated_ptr.is_null() {
69        return core::ptr::null_mut();
70    }
71    layout.store_and_return(allocated_ptr).cast()
72}
73
74unsafe extern "C" fn aligned_free(ptr: *mut core::ffi::c_void) {
75    if ptr.is_null() {
76        return;
77    }
78    let (allocated_ptr, layout) = AlignedLayout::from_raw(ptr);
79    alloc::alloc::dealloc(allocated_ptr.cast(), layout.0);
80}
81
82const MEMORY_HOOKS: ffi::roaring_memory_t = ffi::roaring_memory_t {
83    malloc: Some(malloc),
84    realloc: Some(realloc),
85    calloc: Some(calloc),
86    free: Some(free),
87    aligned_malloc: Some(aligned_malloc),
88    aligned_free: Some(aligned_free),
89};
90
91/// Install custom memory allocation hooks for `CRoaring` which will use rust's global allocator.
92///
93/// # Safety
94///
95/// The caller must ensure there are not any objects allocated by `CRoaring` at the time this
96/// function is called.
97///
98/// Ideally, this function should be called early in the program's execution, before any other
99/// `CRoaring` functions are called.
100pub unsafe fn configure_rust_alloc() {
101    ffi::roaring_init_memory_hook(MEMORY_HOOKS);
102}
103
104#[test]
105fn impossible_aligned_alloc() {
106    unsafe {
107        let ptr = aligned_malloc(usize::MAX, usize::MAX);
108        assert!(ptr.is_null());
109
110        let max_pow_2 = 1usize << (size_of::<usize>() * 8 - 1);
111        let ptr = aligned_malloc(max_pow_2, max_pow_2);
112        assert!(ptr.is_null());
113
114        let max_pow_2_isize = 1usize << (size_of::<usize>() * 8 - 2);
115        let ptr = aligned_malloc(max_pow_2_isize, max_pow_2_isize);
116        assert!(ptr.is_null());
117    }
118}
croaring/rust_alloc/global_alloc.rs

croaring/rust_alloc/
global_alloc.rs
