1 use super::prev_power_of_two;
2 use alloc::collections::BTreeSet;
3 use core::alloc::Layout;
4 use core::cmp::{max, min};
5 use core::ops::Range;
6 
7 #[cfg(feature = "use_spin")]
8 use core::ops::Deref;
9 #[cfg(feature = "use_spin")]
10 use spin::Mutex;
11 
12 /// A frame allocator that uses buddy system, requiring a global allocator.
13 ///
14 /// The max order of the allocator is specified via the const generic parameter `ORDER`. The frame
15 /// allocator will only be able to allocate ranges of size up to 2<sup>ORDER</sup>, out of a total
16 /// range of size at most 2<sup>ORDER + 1</sup> - 1.
17 ///
18 /// # Usage
19 ///
20 /// Create a frame allocator and add some frames to it:
21 /// ```
22 /// use buddy_system_allocator::*;
23 /// let mut frame = FrameAllocator::<32>::new();
24 /// assert!(frame.alloc(1).is_none());
25 ///
26 /// frame.add_frame(0, 3);
27 /// let num = frame.alloc(1);
28 /// assert_eq!(num, Some(2));
29 /// let num = frame.alloc(2);
30 /// assert_eq!(num, Some(0));
31 /// ```
32 pub struct FrameAllocator<const ORDER: usize = 32> {
33     // buddy system with max order of ORDER
34     free_list: [BTreeSet<usize>; ORDER],
35 
36     // statistics
37     allocated: usize,
38     total: usize,
39 }
40 
41 impl<const ORDER: usize> FrameAllocator<ORDER> {
42     /// Create an empty frame allocator
new() -> Self43     pub const fn new() -> Self {
44         Self {
45             free_list: [const { BTreeSet::new() }; ORDER],
46             allocated: 0,
47             total: 0,
48         }
49     }
50 
51     /// Add a range of frame number [start, end) to the allocator
add_frame(&mut self, start: usize, end: usize)52     pub fn add_frame(&mut self, start: usize, end: usize) {
53         assert!(start <= end);
54 
55         let mut total = 0;
56         let mut current_start = start;
57 
58         while current_start < end {
59             let lowbit = if current_start > 0 {
60                 current_start & (!current_start + 1)
61             } else {
62                 32
63             };
64             let size = min(
65                 min(lowbit, prev_power_of_two(end - current_start)),
66                 1 << (ORDER - 1),
67             );
68             total += size;
69 
70             self.free_list[size.trailing_zeros() as usize].insert(current_start);
71             current_start += size;
72         }
73 
74         self.total += total;
75     }
76 
77     /// Add a range of frames to the allocator.
insert(&mut self, range: Range<usize>)78     pub fn insert(&mut self, range: Range<usize>) {
79         self.add_frame(range.start, range.end);
80     }
81 
82     /// Allocate a range of frames from the allocator, returning the first frame of the allocated
83     /// range.
alloc(&mut self, count: usize) -> Option<usize>84     pub fn alloc(&mut self, count: usize) -> Option<usize> {
85         let size = count.next_power_of_two();
86         self.alloc_power_of_two(size)
87     }
88 
89     /// Allocate a range of frames with the given size and alignment from the allocator, returning
90     /// the first frame of the allocated range.
91     /// The allocated size is the maximum of the next power of two of the given size and the
92     /// alignment.
alloc_aligned(&mut self, layout: Layout) -> Option<usize>93     pub fn alloc_aligned(&mut self, layout: Layout) -> Option<usize> {
94         let size = max(layout.size().next_power_of_two(), layout.align());
95         self.alloc_power_of_two(size)
96     }
97 
98     /// Allocate a range of frames of the given size from the allocator. The size must be a power of
99     /// two. The allocated range will have alignment equal to the size.
alloc_power_of_two(&mut self, size: usize) -> Option<usize>100     fn alloc_power_of_two(&mut self, size: usize) -> Option<usize> {
101         let class = size.trailing_zeros() as usize;
102         for i in class..self.free_list.len() {
103             // Find the first non-empty size class
104             if !self.free_list[i].is_empty() {
105                 // Split buffers
106                 for j in (class + 1..i + 1).rev() {
107                     if let Some(block_ref) = self.free_list[j].iter().next() {
108                         let block = *block_ref;
109                         self.free_list[j - 1].insert(block + (1 << (j - 1)));
110                         self.free_list[j - 1].insert(block);
111                         self.free_list[j].remove(&block);
112                     } else {
113                         return None;
114                     }
115                 }
116 
117                 let result = self.free_list[class].iter().next();
118                 if let Some(result_ref) = result {
119                     let result = *result_ref;
120                     self.free_list[class].remove(&result);
121                     self.allocated += size;
122                     return Some(result);
123                 } else {
124                     return None;
125                 }
126             }
127         }
128         None
129     }
130 
131     /// Deallocate a range of frames [frame, frame+count) from the frame allocator.
132     ///
133     /// The range should be exactly the same when it was allocated, as in heap allocator
dealloc(&mut self, start_frame: usize, count: usize)134     pub fn dealloc(&mut self, start_frame: usize, count: usize) {
135         let size = count.next_power_of_two();
136         self.dealloc_power_of_two(start_frame, size)
137     }
138 
139     /// Deallocate a range of frames which was previously allocated by [`alloc_aligned`].
140     ///
141     /// The layout must be exactly the same as when it was allocated.
dealloc_aligned(&mut self, start_frame: usize, layout: Layout)142     pub fn dealloc_aligned(&mut self, start_frame: usize, layout: Layout) {
143         let size = max(layout.size().next_power_of_two(), layout.align());
144         self.dealloc_power_of_two(start_frame, size)
145     }
146 
147     /// Deallocate a range of frames with the given size from the allocator. The size must be a
148     /// power of two.
dealloc_power_of_two(&mut self, start_frame: usize, size: usize)149     fn dealloc_power_of_two(&mut self, start_frame: usize, size: usize) {
150         let class = size.trailing_zeros() as usize;
151 
152         // Merge free buddy lists
153         let mut current_ptr = start_frame;
154         let mut current_class = class;
155         while current_class < self.free_list.len() {
156             let buddy = current_ptr ^ (1 << current_class);
157             if self.free_list[current_class].remove(&buddy) {
158                 // Free buddy found
159                 current_ptr = min(current_ptr, buddy);
160                 current_class += 1;
161             } else {
162                 self.free_list[current_class].insert(current_ptr);
163                 break;
164             }
165         }
166 
167         self.allocated -= size;
168     }
169 }
170 
171 /// A locked version of `FrameAllocator`
172 ///
173 /// # Usage
174 ///
175 /// Create a locked frame allocator and add frames to it:
176 /// ```
177 /// use buddy_system_allocator::*;
178 /// let mut frame = LockedFrameAllocator::<32>::new();
179 /// assert!(frame.lock().alloc(1).is_none());
180 ///
181 /// frame.lock().add_frame(0, 3);
182 /// let num = frame.lock().alloc(1);
183 /// assert_eq!(num, Some(2));
184 /// let num = frame.lock().alloc(2);
185 /// assert_eq!(num, Some(0));
186 /// ```
187 #[cfg(feature = "use_spin")]
188 pub struct LockedFrameAllocator<const ORDER: usize = 32>(Mutex<FrameAllocator<ORDER>>);
189 
190 #[cfg(feature = "use_spin")]
191 impl<const ORDER: usize> LockedFrameAllocator<ORDER> {
192     /// Creates an empty heap
new() -> Self193     pub fn new() -> Self {
194         Self(Mutex::new(FrameAllocator::new()))
195     }
196 }
197 
198 #[cfg(feature = "use_spin")]
199 impl<const ORDER: usize> Deref for LockedFrameAllocator<ORDER> {
200     type Target = Mutex<FrameAllocator<ORDER>>;
201 
deref(&self) -> &Mutex<FrameAllocator<ORDER>>202     fn deref(&self) -> &Mutex<FrameAllocator<ORDER>> {
203         &self.0
204     }
205 }
206