1use parking_lot::RwLock;
7use std::sync::atomic::{AtomicPtr, Ordering};
8
9pub type AllocFn = fn(size: usize, userdata: *mut u8) -> *mut u8;
11pub type ReallocFn = fn(ptr: *mut u8, size: usize, userdata: *mut u8) -> *mut u8;
13pub type FreeFn = fn(ptr: *mut u8, userdata: *mut u8);
15
16struct AllocatorHooks {
17 alloc: AllocFn,
18 realloc: ReallocFn,
19 free: FreeFn,
20 userdata: AtomicPtr<u8>,
21}
22
23impl AllocatorHooks {
24 const fn std_defaults() -> Self {
25 Self {
26 alloc: std_alloc,
27 realloc: std_realloc,
28 free: std_free,
29 userdata: AtomicPtr::new(std::ptr::null_mut()),
30 }
31 }
32
33 fn userdata_ptr(&self) -> *mut u8 {
34 self.userdata.load(Ordering::Acquire)
35 }
36}
37
38static ALLOCATOR: RwLock<AllocatorHooks> = RwLock::new(AllocatorHooks::std_defaults());
39
40fn std_alloc(size: usize, _ud: *mut u8) -> *mut u8 {
41 if size == 0 {
42 return std::ptr::null_mut();
43 }
44 unsafe { libc::malloc(size) as *mut u8 }
45}
46
47fn std_realloc(ptr: *mut u8, size: usize, _ud: *mut u8) -> *mut u8 {
48 if ptr.is_null() {
49 return std_alloc(size, _ud);
50 }
51 if size == 0 {
52 std_free(ptr, _ud);
53 return std::ptr::null_mut();
54 }
55 unsafe { libc::realloc(ptr as *mut libc::c_void, size) as *mut u8 }
56}
57
58fn std_free(ptr: *mut u8, _ud: *mut u8) {
59 if !ptr.is_null() {
60 unsafe { libc::free(ptr as *mut libc::c_void) }
61 }
62}
63
64pub fn set_allocator(alloc: AllocFn, realloc: ReallocFn, free: FreeFn, userdata: *mut u8) {
66 let mut hooks = ALLOCATOR.write();
67 hooks.alloc = alloc;
68 hooks.realloc = realloc;
69 hooks.free = free;
70 hooks.userdata.store(userdata, Ordering::Release);
71}
72
73fn with_hooks<R>(f: impl FnOnce(&AllocatorHooks) -> R) -> R {
74 let hooks = ALLOCATOR.read();
75 f(&hooks)
76}
77
78pub fn lrtmp2_malloc(size: usize) -> *mut u8 {
80 with_hooks(|hooks| (hooks.alloc)(size, hooks.userdata_ptr()))
81}
82
83pub fn lrtmp2_calloc(nmemb: usize, size: usize) -> *mut u8 {
85 if nmemb != 0 && size > usize::MAX / nmemb {
86 return std::ptr::null_mut();
87 }
88 let total = nmemb * size;
89 let p = lrtmp2_malloc(total);
90 if !p.is_null() {
91 unsafe {
92 std::ptr::write_bytes(p, 0, total);
93 }
94 }
95 p
96}
97
98pub fn lrtmp2_realloc(ptr: *mut u8, size: usize) -> *mut u8 {
100 with_hooks(|hooks| (hooks.realloc)(ptr, size, hooks.userdata_ptr()))
101}
102
103pub fn lrtmp2_free(ptr: *mut u8) {
105 if !ptr.is_null() {
106 with_hooks(|hooks| (hooks.free)(ptr, hooks.userdata_ptr()));
107 }
108}
109
110pub fn alloc_vec<T: Copy + Default>(n: usize) -> Vec<T> {
112 vec![T::default(); n]
113}