c/
alloc.rs

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

c/
alloc.rs
