lite_alloc/
lib.rs

1#![no_std]
2
3#[cfg(not(target_arch = "wasm32"))]
4extern crate std;
5
6#[cfg(not(target_arch = "wasm32"))]
7extern crate alloc;
8
9/// WebAssembly memory page count.
10///
11/// WebAssembly 内存页数量。
12#[derive(Eq, PartialEq, Debug, Clone, Copy)]
13pub struct PageCount(pub usize);
14
15impl PageCount {
16    pub fn size_in_bytes(self) -> usize {
17        self.0 * PAGE_SIZE
18    }
19}
20
21/// WebAssembly page size, in bytes (64KB).
22///
23/// WebAssembly 页大小，单位字节（64KB）。
24pub const PAGE_SIZE: usize = 65536;
25
26// Remove MemoryGrower trait, use function directly
27// 移除 trait MemoryGrower，直接写成函数
28#[cfg(target_arch = "wasm32")]
29#[inline(always)]
30pub unsafe fn grow_memory(pages: usize) -> usize {
31    core::arch::wasm32::memory_grow(0, pages)
32}
33
34#[cfg(not(target_arch = "wasm32"))]
35mod host_memory {
36    use super::PAGE_SIZE;
37    use std::alloc::{Layout, alloc, dealloc};
38    use std::cell::RefCell;
39    use std::ptr;
40
41    // Simulate 128MB of addressable WASM memory space per thread
42    // 模拟每线程 128MB 的可寻址 WASM 内存空间
43    const MOCK_MEMORY_SIZE: usize = 128 * 1024 * 1024;
44
45    struct MockMemory {
46        base_ptr: *mut u8,
47        current_pages: usize,
48    }
49
50    impl MockMemory {
51        fn new() -> Self {
52            unsafe {
53                let layout = Layout::from_size_align(MOCK_MEMORY_SIZE, PAGE_SIZE).unwrap();
54                let ptr = alloc(layout);
55                if ptr.is_null() {
56                    // Panic immediately if we can't allocate the mock heap
57                    // 如果无法分配模拟堆，立即 Panic
58                    panic!("Failed to allocate mock WASM memory");
59                }
60                // Initialize memory to zero, similar to WASM behavior
61                // 将内存初始化为零，类似于 WASM 行为
62                ptr::write_bytes(ptr, 0, MOCK_MEMORY_SIZE);
63                Self {
64                    base_ptr: ptr,
65                    current_pages: 0,
66                }
67            }
68        }
69    }
70
71    impl Drop for MockMemory {
72        fn drop(&mut self) {
73            unsafe {
74                let layout = Layout::from_size_align(MOCK_MEMORY_SIZE, PAGE_SIZE).unwrap();
75                dealloc(self.base_ptr, layout);
76            }
77        }
78    }
79
80    // Thread local storage for the mock memory
81    // 模拟内存的线程局部存储
82    std::thread_local! {
83        static MEMORY: RefCell<MockMemory> = RefCell::new(MockMemory::new());
84    }
85
86    pub unsafe fn grow_memory_impl(pages: usize) -> usize {
87        MEMORY.with(|mem| {
88            let mut mem = mem.borrow_mut();
89
90            // Check if we have enough space in our pre-allocated buffer
91            // 检查预分配缓冲区中是否有足够的空间
92            if (mem.current_pages + pages) * PAGE_SIZE > MOCK_MEMORY_SIZE {
93                return usize::MAX;
94            }
95
96            // Calculate the start of the new memory block as an absolute address
97            // 计算新内存块的起始绝对地址
98            let start_addr = mem.base_ptr as usize + mem.current_pages * PAGE_SIZE;
99
100            // Since our allocators expect the return value to be (Address / PAGE_SIZE),
101            // and they will reconstruct the address by (RetVal * PAGE_SIZE),
102            // we must return the absolute page index.
103            // 由于我们的分配器期望返回值为 (Address / PAGE_SIZE)，
104            // 并且它们将通过 (RetVal * PAGE_SIZE) 重建地址，
105            // 我们必须返回绝对页索引。
106            let ret_page_index = start_addr / PAGE_SIZE;
107
108            // Advance the usage counter
109            // 增加使用计数
110            mem.current_pages += pages;
111
112            // Zero out the newly allocated pages (emulate WASM grow behavior)
113            // Note: We zeroed the whole block on init, but reset_memory might not zero everything
114            // if we optimize it. Let's ensure it's zeroed here if we change reset policy.
115            // For now, init zeros everything and reset zeros used, so it's fine.
116            // 将新分配的页面清零（模拟 WASM 增长行为）
117            // 注意：我们在初始化时将整个块清零，但如果进行优化，reset_memory 可能不会清零所有内容。
118            // 如果我们要更改重置策略，请确保在此处清零。
119            // 目前，初始化会清零所有内容，重置会清零已使用的内容，所以没问题。
120
121            ret_page_index
122        })
123    }
124
125    pub unsafe fn reset_memory() {
126        MEMORY.with(|mem| {
127            let mut mem = mem.borrow_mut();
128            // Zero out the used memory so the next test starts clean
129            // 将已使用的内存清零，以便下一个测试从干净的状态开始
130            let used_bytes = mem.current_pages * PAGE_SIZE;
131            if used_bytes > 0 {
132                unsafe {
133                    ptr::write_bytes(mem.base_ptr, 0, used_bytes);
134                }
135            }
136            mem.current_pages = 0;
137        });
138    }
139}
140
141#[cfg(not(target_arch = "wasm32"))]
142pub unsafe fn grow_memory(pages: usize) -> usize {
143    unsafe { host_memory::grow_memory_impl(pages) }
144}
145
146/// For Test/Bench only: Reset the mock heap memory of the current thread
147///
148/// 仅用于测试/Bench：重置当前线程的模拟堆内存
149#[cfg(not(target_arch = "wasm32"))]
150pub fn reset_heap() {
151    unsafe {
152        host_memory::reset_memory();
153    }
154}
155
156pub mod single_threaded {
157    mod bump_freelist;
158    mod freelist;
159    mod segregated_bump;
160
161    pub use bump_freelist::BumpFreeListAllocator;
162    pub use freelist::FreeListAllocator;
163    pub use segregated_bump::SegregatedBumpAllocator;
164}
lite_alloc/lib.rs

lite_alloc/
lib.rs
