c/
alloc.rs

1// Allocator: segregated free lists for small blocks, mmap for large blocks.
2//
3// This replaces the old mmap-per-allocation allocator. Small allocations are
4// carved from page-aligned chunks and recycled through per-size free lists;
5// large allocations still use mmap directly. A single short-lived spin lock
6// protects the free lists. fork() resets it in the child.
7
8const ALLOC_STEP: usize = 16;
9const ALLOC_MIN_CLASS: usize = 32;
10const ALLOC_MAX_CLASS: usize = 32768;
11const ALLOC_CLASSES: usize = ALLOC_MAX_CLASS / ALLOC_STEP + 1;
12const ALLOC_HDR: usize = core::mem::size_of::<AllocBlock>();
13const ALLOC_LARGE: usize = 1usize << 63;
14const ALLOC_ALIGNED: usize = 1usize << 62;
15const ALLOC_PAGES_MASK: usize = ALLOC_LARGE - 1;
16
17#[repr(C)]
18struct AllocBlock {
19    next: *mut AllocBlock,
20    tag: usize,
21}
22
23static ALLOC_LOCK: AtomicUsize = AtomicUsize::new(0);
24static mut ALLOC_FREE: [*mut AllocBlock; ALLOC_CLASSES] =
25    [null_mut(); ALLOC_CLASSES];
26
27#[inline]
28unsafe fn alloc_lock() {
29    while ALLOC_LOCK.swap(1, Ordering::Acquire) != 0 {
30        core::hint::spin_loop();
31    }
32}
33
34#[inline]
35unsafe fn alloc_unlock() {
36    ALLOC_LOCK.store(0, Ordering::Release);
37}
38
39pub(crate) unsafe fn allocator_after_fork() {
40    ALLOC_LOCK.store(0, Ordering::Release);
41}
42
43#[inline]
44fn alloc_class_for(size: usize) -> Option<usize> {
45    let total = size.checked_add(ALLOC_HDR)?;
46    let class = if total <= ALLOC_MIN_CLASS {
47        ALLOC_MIN_CLASS
48    } else {
49        match total.checked_add(ALLOC_STEP - 1) {
50            Some(v) => v & !(ALLOC_STEP - 1),
51            None => return None,
52        }
53    };
54    if class <= ALLOC_MAX_CLASS { Some(class) } else { None }
55}
56
57unsafe fn alloc_mmap(size: usize) -> *mut u8 {
58    let pages = (size + PAGE - 1) & !(PAGE - 1);
59    let p = sys_mmap(
60        null_mut(),
61        pages,
62        PROT_READ | PROT_WRITE,
63        MAP_PRIVATE | MAP_ANONYMOUS,
64        -1,
65        0,
66    );
67    if p == MMAP_FAILED { null_mut() } else { p }
68}
69
70unsafe fn alloc_large(size: usize) -> *mut u8 {
71    let total = match size.checked_add(ALLOC_HDR) {
72        Some(v) => v,
73        None => {
74            ERRNO!() = ENOMEM;
75            return null_mut();
76        }
77    };
78    if total > usize::MAX - (PAGE - 1) {
79        ERRNO!() = ENOMEM;
80        return null_mut();
81    }
82    let pages = (total + PAGE - 1) / PAGE;
83    let map_size = pages * PAGE;
84    let base = alloc_mmap(map_size);
85    if base.is_null() {
86        return null_mut();
87    }
88    let block = base as *mut AllocBlock;
89    (*block).next = block;
90    (*block).tag = ALLOC_LARGE | pages;
91    base.add(ALLOC_HDR)
92}
93
94unsafe fn alloc_aligned_impl(alignment: usize, size: usize) -> *mut u8 {
95    if alignment == 0 || !alignment.is_power_of_two() {
96        ERRNO!() = EINVAL;
97        return null_mut();
98    }
99    if alignment <= ALLOC_STEP {
100        return malloc(size) as *mut u8;
101    }
102    let total = match size.checked_add(alignment).and_then(|v| v.checked_add(ALLOC_HDR)) {
103        Some(v) => v,
104        None => {
105            ERRNO!() = ENOMEM;
106            return null_mut();
107        }
108    };
109    if total > usize::MAX - (PAGE - 1) {
110        ERRNO!() = ENOMEM;
111        return null_mut();
112    }
113    let pages = (total + PAGE - 1) / PAGE;
114    let map_size = pages * PAGE;
115    let base = alloc_mmap(map_size);
116    if base.is_null() {
117        return null_mut();
118    }
119    let data_addr = (base as usize + ALLOC_HDR + alignment - 1) & !(alignment - 1);
120    let block = (data_addr - ALLOC_HDR) as *mut AllocBlock;
121    (*block).next = base as *mut AllocBlock;
122    (*block).tag = ALLOC_ALIGNED | pages;
123    data_addr as *mut u8
124}
125
126pub(crate) unsafe fn alloc_usable_size(ptr: *mut u8) -> usize {
127    if ptr.is_null() {
128        return 0;
129    }
130    let block = ptr.sub(ALLOC_HDR) as *mut AllocBlock;
131    let tag = (*block).tag;
132    if tag & ALLOC_ALIGNED != 0 {
133        let base = (*block).next as usize;
134        let pages = tag & ALLOC_PAGES_MASK;
135        pages * PAGE - (ptr as usize - base)
136    } else if tag & ALLOC_LARGE != 0 {
137        let pages = tag & ALLOC_PAGES_MASK;
138        pages * PAGE - ALLOC_HDR
139    } else if tag >= ALLOC_MIN_CLASS && tag <= ALLOC_MAX_CLASS {
140        tag - ALLOC_HDR
141    } else {
142        0
143    }
144}
145
146#[no_mangle]
147pub unsafe extern "C" fn malloc(size: SizeT) -> *mut c_void {
148    let class = match alloc_class_for(size) {
149        Some(c) => c,
150        None => return alloc_large(size) as *mut c_void,
151    };
152    let idx = class / ALLOC_STEP;
153
154    alloc_lock();
155    let head = ALLOC_FREE[idx];
156    if !head.is_null() {
157        ALLOC_FREE[idx] = (*head).next;
158        alloc_unlock();
159        return (head as *mut u8).add(ALLOC_HDR) as *mut c_void;
160    }
161
162    // Refill this size class with a chunk of blocks.
163    let chunk_size = if class < 4096 { 65536 } else { class * 2 };
164    let map_size = (chunk_size + PAGE - 1) & !(PAGE - 1);
165    let base = alloc_mmap(map_size);
166    if base.is_null() {
167        alloc_unlock();
168        return null_mut();
169    }
170    let count = map_size / class;
171    let mut i = count;
172    while i > 1 {
173        i -= 1;
174        let block = base.add(i * class) as *mut AllocBlock;
175        (*block).tag = class;
176        (*block).next = ALLOC_FREE[idx];
177        ALLOC_FREE[idx] = block;
178    }
179    let first = base as *mut AllocBlock;
180    (*first).tag = class;
181    alloc_unlock();
182    base.add(ALLOC_HDR) as *mut c_void
183}
184
185#[no_mangle]
186pub unsafe extern "C" fn free(ptr: *mut c_void) {
187    if ptr.is_null() {
188        return;
189    }
190    let block = (ptr as *mut u8).sub(ALLOC_HDR) as *mut AllocBlock;
191    let tag = (*block).tag;
192    if tag & ALLOC_ALIGNED != 0 {
193        let base = (*block).next as *mut u8;
194        let pages = tag & ALLOC_PAGES_MASK;
195        sys_munmap(base, pages * PAGE);
196        return;
197    }
198    if tag & ALLOC_LARGE != 0 {
199        let pages = tag & ALLOC_PAGES_MASK;
200        sys_munmap(block as *mut u8, pages * PAGE);
201        return;
202    }
203    if tag < ALLOC_MIN_CLASS || tag > ALLOC_MAX_CLASS {
204        return;
205    }
206    let idx = tag / ALLOC_STEP;
207    alloc_lock();
208    (*block).next = ALLOC_FREE[idx];
209    ALLOC_FREE[idx] = block;
210    alloc_unlock();
211}
212
213#[no_mangle]
214pub unsafe extern "C" fn calloc(count: SizeT, size: SizeT) -> *mut c_void {
215    let total = match count.checked_mul(size) {
216        Some(v) => v,
217        None => {
218            ERRNO!() = ENOMEM;
219            return null_mut();
220        }
221    };
222    let ptr = malloc(total);
223    if !ptr.is_null() {
224        core::ptr::write_bytes(ptr as *mut u8, 0, total);
225    }
226    ptr
227}
228
229#[no_mangle]
230pub unsafe extern "C" fn realloc(ptr: *mut c_void, new_size: SizeT) -> *mut c_void {
231    if ptr.is_null() {
232        return malloc(new_size);
233    }
234    if new_size == 0 {
235        free(ptr);
236        return null_mut();
237    }
238    let old_size = alloc_usable_size(ptr as *mut u8);
239    if new_size <= old_size {
240        return ptr;
241    }
242    let new_ptr = malloc(new_size);
243    if new_ptr.is_null() {
244        return null_mut();
245    }
246    core::ptr::copy_nonoverlapping(ptr as *const u8, new_ptr as *mut u8, old_size);
247    free(ptr);
248    new_ptr
249}
c/alloc.rs

c/
alloc.rs
