Skip to main content

librtmp2/
alloc.rs

1//! Custom allocator hooks for librtmp2
2//!
3//! Mirrors `src/core/alloc.h` and `src/core/alloc.c`.
4//! Provides pluggable allocation functions with standard defaults.
5
6use parking_lot::RwLock;
7use std::sync::atomic::{AtomicPtr, Ordering};
8
9/// Allocator function type.
10pub type AllocFn = fn(size: usize, userdata: *mut u8) -> *mut u8;
11/// Reallocator function type.
12pub type ReallocFn = fn(ptr: *mut u8, size: usize, userdata: *mut u8) -> *mut u8;
13/// Free function type.
14pub 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
64/// Set custom allocator functions.
65pub 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
78/// Allocate `size` bytes using the current allocator.
79pub fn lrtmp2_malloc(size: usize) -> *mut u8 {
80    with_hooks(|hooks| (hooks.alloc)(size, hooks.userdata_ptr()))
81}
82
83/// Allocate zeroed memory for `nmemb` elements of `size` bytes each.
84pub 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
98/// Reallocate memory.
99pub fn lrtmp2_realloc(ptr: *mut u8, size: usize) -> *mut u8 {
100    with_hooks(|hooks| (hooks.realloc)(ptr, size, hooks.userdata_ptr()))
101}
102
103/// Free memory.
104pub fn lrtmp2_free(ptr: *mut u8) {
105    if !ptr.is_null() {
106        with_hooks(|hooks| (hooks.free)(ptr, hooks.userdata_ptr()));
107    }
108}
109
110/// Allocate a Vec-backed buffer (idiomatic Rust helper).
111pub fn alloc_vec<T: Copy + Default>(n: usize) -> Vec<T> {
112    vec![T::default(); n]
113}