axalloc/
lib.rs

1//! [ArceOS](https://github.com/arceos-org/arceos) global memory allocator.
2//!
3//! It provides [`GlobalAllocator`], which implements the trait
4//! [`core::alloc::GlobalAlloc`]. A static global variable of type
5//! [`GlobalAllocator`] is defined with the `#[global_allocator]` attribute, to
6//! be registered as the standard library’s default allocator.
7
8#![no_std]
9
10#[macro_use]
11extern crate log;
12extern crate alloc;
13
14mod page;
15
16use allocator::{AllocResult, BaseAllocator, BitmapPageAllocator, ByteAllocator, PageAllocator};
17use core::alloc::{GlobalAlloc, Layout};
18use core::ptr::NonNull;
19use kspin::SpinNoIrq;
20
21const PAGE_SIZE: usize = 0x1000;
22const MIN_HEAP_SIZE: usize = 0x8000; // 32 K
23
24pub use page::GlobalPage;
25
26cfg_if::cfg_if! {
27    if #[cfg(feature = "slab")] {
28        /// The default byte allocator.
29        pub type DefaultByteAllocator = allocator::SlabByteAllocator;
30    } else if #[cfg(feature = "buddy")] {
31        /// The default byte allocator.
32        pub type DefaultByteAllocator = allocator::BuddyByteAllocator;
33    } else if #[cfg(feature = "tlsf")] {
34        /// The default byte allocator.
35        pub type DefaultByteAllocator = allocator::TlsfByteAllocator;
36    }
37}
38
39/// The global allocator used by ArceOS.
40///
41/// It combines a [`ByteAllocator`] and a [`PageAllocator`] into a simple
42/// two-level allocator: firstly tries allocate from the byte allocator, if
43/// there is no memory, asks the page allocator for more memory and adds it to
44/// the byte allocator.
45///
46/// Currently, [`TlsfByteAllocator`] is used as the byte allocator, while
47/// [`BitmapPageAllocator`] is used as the page allocator.
48///
49/// [`TlsfByteAllocator`]: allocator::TlsfByteAllocator
50pub struct GlobalAllocator {
51    balloc: SpinNoIrq<DefaultByteAllocator>,
52    palloc: SpinNoIrq<BitmapPageAllocator<PAGE_SIZE>>,
53}
54
55impl GlobalAllocator {
56    /// Creates an empty [`GlobalAllocator`].
57    pub const fn new() -> Self {
58        Self {
59            balloc: SpinNoIrq::new(DefaultByteAllocator::new()),
60            palloc: SpinNoIrq::new(BitmapPageAllocator::new()),
61        }
62    }
63
64    /// Returns the name of the allocator.
65    pub const fn name(&self) -> &'static str {
66        cfg_if::cfg_if! {
67            if #[cfg(feature = "slab")] {
68                "slab"
69            } else if #[cfg(feature = "buddy")] {
70                "buddy"
71            } else if #[cfg(feature = "tlsf")] {
72                "TLSF"
73            }
74        }
75    }
76
77    /// Initializes the allocator with the given region.
78    ///
79    /// It firstly adds the whole region to the page allocator, then allocates
80    /// a small region (32 KB) to initialize the byte allocator. Therefore,
81    /// the given region must be larger than 32 KB.
82    pub fn init(&self, start_vaddr: usize, size: usize) {
83        assert!(size > MIN_HEAP_SIZE);
84        let init_heap_size = MIN_HEAP_SIZE;
85        self.palloc.lock().init(start_vaddr, size);
86        let heap_ptr = self
87            .alloc_pages(init_heap_size / PAGE_SIZE, PAGE_SIZE)
88            .unwrap();
89        self.balloc.lock().init(heap_ptr, init_heap_size);
90    }
91
92    /// Add the given region to the allocator.
93    ///
94    /// It will add the whole region to the byte allocator.
95    pub fn add_memory(&self, start_vaddr: usize, size: usize) -> AllocResult {
96        self.balloc.lock().add_memory(start_vaddr, size)
97    }
98
99    /// Allocate arbitrary number of bytes. Returns the left bound of the
100    /// allocated region.
101    ///
102    /// It firstly tries to allocate from the byte allocator. If there is no
103    /// memory, it asks the page allocator for more memory and adds it to the
104    /// byte allocator.
105    pub fn alloc(&self, layout: Layout) -> AllocResult<NonNull<u8>> {
106        // simple two-level allocator: if no heap memory, allocate from the page allocator.
107        let mut balloc = self.balloc.lock();
108        loop {
109            if let Ok(ptr) = balloc.alloc(layout) {
110                return Ok(ptr);
111            } else {
112                let old_size = balloc.total_bytes();
113                let expand_size = old_size
114                    .max(layout.size())
115                    .next_power_of_two()
116                    .max(PAGE_SIZE);
117                let heap_ptr = self.alloc_pages(expand_size / PAGE_SIZE, PAGE_SIZE)?;
118                debug!(
119                    "expand heap memory: [{:#x}, {:#x})",
120                    heap_ptr,
121                    heap_ptr + expand_size
122                );
123                balloc.add_memory(heap_ptr, expand_size)?;
124            }
125        }
126    }
127
128    /// Gives back the allocated region to the byte allocator.
129    ///
130    /// The region should be allocated by [`alloc`], and `align_pow2` should be
131    /// the same as the one used in [`alloc`]. Otherwise, the behavior is
132    /// undefined.
133    ///
134    /// [`alloc`]: GlobalAllocator::alloc
135    pub fn dealloc(&self, pos: NonNull<u8>, layout: Layout) {
136        self.balloc.lock().dealloc(pos, layout)
137    }
138
139    /// Allocates contiguous pages.
140    ///
141    /// It allocates `num_pages` pages from the page allocator.
142    ///
143    /// `align_pow2` must be a power of 2, and the returned region bound will be
144    /// aligned to it.
145    pub fn alloc_pages(&self, num_pages: usize, align_pow2: usize) -> AllocResult<usize> {
146        self.palloc.lock().alloc_pages(num_pages, align_pow2)
147    }
148
149    /// Gives back the allocated pages starts from `pos` to the page allocator.
150    ///
151    /// The pages should be allocated by [`alloc_pages`], and `align_pow2`
152    /// should be the same as the one used in [`alloc_pages`]. Otherwise, the
153    /// behavior is undefined.
154    ///
155    /// [`alloc_pages`]: GlobalAllocator::alloc_pages
156    pub fn dealloc_pages(&self, pos: usize, num_pages: usize) {
157        self.palloc.lock().dealloc_pages(pos, num_pages)
158    }
159
160    /// Returns the number of allocated bytes in the byte allocator.
161    pub fn used_bytes(&self) -> usize {
162        self.balloc.lock().used_bytes()
163    }
164
165    /// Returns the number of available bytes in the byte allocator.
166    pub fn available_bytes(&self) -> usize {
167        self.balloc.lock().available_bytes()
168    }
169
170    /// Returns the number of allocated pages in the page allocator.
171    pub fn used_pages(&self) -> usize {
172        self.palloc.lock().used_pages()
173    }
174
175    /// Returns the number of available pages in the page allocator.
176    pub fn available_pages(&self) -> usize {
177        self.palloc.lock().available_pages()
178    }
179}
180
181unsafe impl GlobalAlloc for GlobalAllocator {
182    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
183        if let Ok(ptr) = GlobalAllocator::alloc(self, layout) {
184            ptr.as_ptr()
185        } else {
186            alloc::alloc::handle_alloc_error(layout)
187        }
188    }
189
190    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
191        GlobalAllocator::dealloc(self, NonNull::new(ptr).expect("dealloc null ptr"), layout)
192    }
193}
194
195#[cfg_attr(all(target_os = "none", not(test)), global_allocator)]
196static GLOBAL_ALLOCATOR: GlobalAllocator = GlobalAllocator::new();
197
198/// Returns the reference to the global allocator.
199pub fn global_allocator() -> &'static GlobalAllocator {
200    &GLOBAL_ALLOCATOR
201}
202
203/// Initializes the global allocator with the given memory region.
204///
205/// Note that the memory region bounds are just numbers, and the allocator
206/// does not actually access the region. Users should ensure that the region
207/// is valid and not being used by others, so that the allocated memory is also
208/// valid.
209///
210/// This function should be called only once, and before any allocation.
211pub fn global_init(start_vaddr: usize, size: usize) {
212    debug!(
213        "initialize global allocator at: [{:#x}, {:#x})",
214        start_vaddr,
215        start_vaddr + size
216    );
217    GLOBAL_ALLOCATOR.init(start_vaddr, size);
218}
219
220/// Add the given memory region to the global allocator.
221///
222/// Users should ensure that the region is valid and not being used by others,
223/// so that the allocated memory is also valid.
224///
225/// It's similar to [`global_init`], but can be called multiple times.
226pub fn global_add_memory(start_vaddr: usize, size: usize) -> AllocResult {
227    debug!(
228        "add a memory region to global allocator: [{:#x}, {:#x})",
229        start_vaddr,
230        start_vaddr + size
231    );
232    GLOBAL_ALLOCATOR.add_memory(start_vaddr, size)
233}