range_alloc_arceos/
lib.rs

1#![no_std]
2
3extern crate alloc;
4
5use alloc::vec;
6use alloc::vec::Vec;
7use core::{
8    fmt::Debug,
9    iter::Sum,
10    ops::{Add, AddAssign, Range, Sub},
11};
12
13#[derive(Debug)]
14pub struct RangeAllocator<T> {
15    /// The range this allocator covers.
16    initial_range: Range<T>,
17    /// A Vec of ranges in this heap which are unused.
18    /// Must be ordered with ascending range start to permit short circuiting allocation.
19    /// No two ranges in this vec may overlap.
20    free_ranges: Vec<Range<T>>,
21}
22
23#[derive(Clone, Debug, PartialEq)]
24pub struct RangeAllocationError<T> {
25    pub fragmented_free_length: T,
26}
27
28impl<T> RangeAllocator<T>
29where
30    T: Clone + Copy + Add<Output = T> + AddAssign + Sub<Output = T> + Eq + PartialOrd + Debug,
31{
32    pub fn new(range: Range<T>) -> Self {
33        RangeAllocator {
34            initial_range: range.clone(),
35            free_ranges: vec![range],
36        }
37    }
38
39    pub fn initial_range(&self) -> &Range<T> {
40        &self.initial_range
41    }
42
43    pub fn grow_to(&mut self, new_end: T) {
44        let initial_range_end = self.initial_range.end;
45        if let Some(last_range) = self
46            .free_ranges
47            .last_mut()
48            .filter(|last_range| last_range.end == initial_range_end)
49        {
50            last_range.end = new_end;
51        } else {
52            self.free_ranges.push(self.initial_range.end..new_end);
53        }
54
55        self.initial_range.end = new_end;
56    }
57
58    pub fn allocate_range(&mut self, length: T) -> Result<Range<T>, RangeAllocationError<T>> {
59        assert_ne!(length + length, length);
60        let mut best_fit: Option<(usize, Range<T>)> = None;
61
62        // This is actually correct. With the trait bound as it is, we have
63        // no way to summon a value of 0 directly, so we make one by subtracting
64        // something from itself. Once the trait bound can be changed, this can
65        // be fixed.
66        #[allow(clippy::eq_op)]
67        let mut fragmented_free_length = length - length;
68        for (index, range) in self.free_ranges.iter().cloned().enumerate() {
69            let range_length = range.end - range.start;
70            fragmented_free_length += range_length;
71            if range_length < length {
72                continue;
73            } else if range_length == length {
74                // Found a perfect fit, so stop looking.
75                best_fit = Some((index, range));
76                break;
77            }
78            best_fit = Some(match best_fit {
79                Some((best_index, best_range)) => {
80                    // Find best fit for this allocation to reduce memory fragmentation.
81                    if range_length < best_range.end - best_range.start {
82                        (index, range)
83                    } else {
84                        (best_index, best_range.clone())
85                    }
86                }
87                None => (index, range),
88            });
89        }
90        match best_fit {
91            Some((index, range)) => {
92                if range.end - range.start == length {
93                    self.free_ranges.remove(index);
94                } else {
95                    self.free_ranges[index].start += length;
96                }
97                Ok(range.start..(range.start + length))
98            }
99            None => Err(RangeAllocationError {
100                fragmented_free_length,
101            }),
102        }
103    }
104
105    pub fn free_range(&mut self, range: Range<T>) {
106        assert!(self.initial_range.start <= range.start && range.end <= self.initial_range.end);
107        assert!(range.start < range.end);
108
109        // Get insertion position.
110        let i = self
111            .free_ranges
112            .iter()
113            .position(|r| r.start > range.start)
114            .unwrap_or(self.free_ranges.len());
115
116        // Try merging with neighboring ranges in the free list.
117        // Before: |left|-(range)-|right|
118        if i > 0 && range.start == self.free_ranges[i - 1].end {
119            // Merge with |left|.
120            self.free_ranges[i - 1].end =
121                if i < self.free_ranges.len() && range.end == self.free_ranges[i].start {
122                    // Check for possible merge with |left| and |right|.
123                    let right = self.free_ranges.remove(i);
124                    right.end
125                } else {
126                    range.end
127                };
128
129            return;
130        } else if i < self.free_ranges.len() && range.end == self.free_ranges[i].start {
131            // Merge with |right|.
132            self.free_ranges[i].start = if i > 0 && range.start == self.free_ranges[i - 1].end {
133                // Check for possible merge with |left| and |right|.
134                let left = self.free_ranges.remove(i - 1);
135                left.start
136            } else {
137                range.start
138            };
139
140            return;
141        }
142
143        // Debug checks
144        assert!(
145            (i == 0 || self.free_ranges[i - 1].end < range.start)
146                && (i >= self.free_ranges.len() || range.end < self.free_ranges[i].start)
147        );
148
149        self.free_ranges.insert(i, range);
150    }
151
152    /// Returns an iterator over allocated non-empty ranges
153    pub fn allocated_ranges(&self) -> impl Iterator<Item = Range<T>> + '_ {
154        let first = match self.free_ranges.first() {
155            Some(Range { ref start, .. }) if *start > self.initial_range.start => {
156                Some(self.initial_range.start..*start)
157            }
158            None => Some(self.initial_range.clone()),
159            _ => None,
160        };
161
162        let last = match self.free_ranges.last() {
163            Some(Range { end, .. }) if *end < self.initial_range.end => {
164                Some(*end..self.initial_range.end)
165            }
166            _ => None,
167        };
168
169        let mid = self
170            .free_ranges
171            .iter()
172            .zip(self.free_ranges.iter().skip(1))
173            .map(|(ra, rb)| ra.end..rb.start);
174
175        first.into_iter().chain(mid).chain(last)
176    }
177
178    pub fn reset(&mut self) {
179        self.free_ranges.clear();
180        self.free_ranges.push(self.initial_range.clone());
181    }
182
183    pub fn is_empty(&self) -> bool {
184        self.free_ranges.len() == 1 && self.free_ranges[0] == self.initial_range
185    }
186}
187
188impl<T: Copy + Sub<Output = T> + Sum> RangeAllocator<T> {
189    pub fn total_available(&self) -> T {
190        self.free_ranges
191            .iter()
192            .map(|range| range.end - range.start)
193            .sum()
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn test_basic_allocation() {
203        let mut alloc = RangeAllocator::new(0..10);
204        // Test if an allocation works
205        assert_eq!(alloc.allocate_range(4), Ok(0..4));
206        assert!(alloc.allocated_ranges().eq(core::iter::once(0..4)));
207        // Free the prior allocation
208        alloc.free_range(0..4);
209        // Make sure the free actually worked
210        assert_eq!(alloc.free_ranges, vec![0..10]);
211        assert!(alloc.allocated_ranges().eq(core::iter::empty()));
212    }
213
214    #[test]
215    fn test_out_of_space() {
216        let mut alloc = RangeAllocator::new(0..10);
217        // Test if the allocator runs out of space correctly
218        assert_eq!(alloc.allocate_range(10), Ok(0..10));
219        assert!(alloc.allocated_ranges().eq(core::iter::once(0..10)));
220        assert!(alloc.allocate_range(4).is_err());
221        alloc.free_range(0..10);
222    }
223
224    #[test]
225    fn test_grow() {
226        let mut alloc = RangeAllocator::new(0..11);
227        // Test if the allocator runs out of space correctly
228        assert_eq!(alloc.allocate_range(10), Ok(0..10));
229        assert!(alloc.allocated_ranges().eq(core::iter::once(0..10)));
230        assert!(alloc.allocate_range(4).is_err());
231        alloc.grow_to(20);
232        assert_eq!(alloc.allocate_range(4), Ok(10..14));
233        alloc.free_range(0..14);
234    }
235
236    #[test]
237    fn test_grow_with_hole_at_start() {
238        let mut alloc = RangeAllocator::new(0..6);
239
240        assert_eq!(alloc.allocate_range(3), Ok(0..3));
241        assert_eq!(alloc.allocate_range(3), Ok(3..6));
242        alloc.free_range(0..3);
243
244        alloc.grow_to(9);
245        assert_eq!(alloc.allocated_ranges().collect::<Vec<_>>(), [3..6]);
246    }
247    #[test]
248    fn test_grow_with_hole_in_middle() {
249        let mut alloc = RangeAllocator::new(0..6);
250
251        assert_eq!(alloc.allocate_range(2), Ok(0..2));
252        assert_eq!(alloc.allocate_range(2), Ok(2..4));
253        assert_eq!(alloc.allocate_range(2), Ok(4..6));
254        alloc.free_range(2..4);
255
256        alloc.grow_to(9);
257        assert_eq!(alloc.allocated_ranges().collect::<Vec<_>>(), [0..2, 4..6]);
258    }
259
260    #[test]
261    fn test_dont_use_block_that_is_too_small() {
262        let mut alloc = RangeAllocator::new(0..10);
263        // Allocate three blocks then free the middle one and check for correct state
264        assert_eq!(alloc.allocate_range(3), Ok(0..3));
265        assert_eq!(alloc.allocate_range(3), Ok(3..6));
266        assert_eq!(alloc.allocate_range(3), Ok(6..9));
267        alloc.free_range(3..6);
268        assert_eq!(alloc.free_ranges, vec![3..6, 9..10]);
269        assert_eq!(
270            alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
271            vec![0..3, 6..9]
272        );
273        // Now request space that the middle block can fill, but the end one can't.
274        assert_eq!(alloc.allocate_range(3), Ok(3..6));
275    }
276
277    #[test]
278    fn test_free_blocks_in_middle() {
279        let mut alloc = RangeAllocator::new(0..100);
280        // Allocate many blocks then free every other block.
281        assert_eq!(alloc.allocate_range(10), Ok(0..10));
282        assert_eq!(alloc.allocate_range(10), Ok(10..20));
283        assert_eq!(alloc.allocate_range(10), Ok(20..30));
284        assert_eq!(alloc.allocate_range(10), Ok(30..40));
285        assert_eq!(alloc.allocate_range(10), Ok(40..50));
286        assert_eq!(alloc.allocate_range(10), Ok(50..60));
287        assert_eq!(alloc.allocate_range(10), Ok(60..70));
288        assert_eq!(alloc.allocate_range(10), Ok(70..80));
289        assert_eq!(alloc.allocate_range(10), Ok(80..90));
290        assert_eq!(alloc.allocate_range(10), Ok(90..100));
291        assert_eq!(alloc.free_ranges, vec![]);
292        assert!(alloc.allocated_ranges().eq(core::iter::once(0..100)));
293        alloc.free_range(10..20);
294        alloc.free_range(30..40);
295        alloc.free_range(50..60);
296        alloc.free_range(70..80);
297        alloc.free_range(90..100);
298        // Check that the right blocks were freed.
299        assert_eq!(
300            alloc.free_ranges,
301            vec![10..20, 30..40, 50..60, 70..80, 90..100]
302        );
303        assert_eq!(
304            alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
305            vec![0..10, 20..30, 40..50, 60..70, 80..90]
306        );
307        // Fragment the memory on purpose a bit.
308        assert_eq!(alloc.allocate_range(6), Ok(10..16));
309        assert_eq!(alloc.allocate_range(6), Ok(30..36));
310        assert_eq!(alloc.allocate_range(6), Ok(50..56));
311        assert_eq!(alloc.allocate_range(6), Ok(70..76));
312        assert_eq!(alloc.allocate_range(6), Ok(90..96));
313        // Check for fragmentation.
314        assert_eq!(
315            alloc.free_ranges,
316            vec![16..20, 36..40, 56..60, 76..80, 96..100]
317        );
318        assert_eq!(
319            alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
320            vec![0..16, 20..36, 40..56, 60..76, 80..96]
321        );
322        // Fill up the fragmentation
323        assert_eq!(alloc.allocate_range(4), Ok(16..20));
324        assert_eq!(alloc.allocate_range(4), Ok(36..40));
325        assert_eq!(alloc.allocate_range(4), Ok(56..60));
326        assert_eq!(alloc.allocate_range(4), Ok(76..80));
327        assert_eq!(alloc.allocate_range(4), Ok(96..100));
328        // Check that nothing is free.
329        assert_eq!(alloc.free_ranges, vec![]);
330        assert!(alloc.allocated_ranges().eq(core::iter::once(0..100)));
331    }
332
333    #[test]
334    fn test_ignore_block_if_another_fits_better() {
335        let mut alloc = RangeAllocator::new(0..10);
336        // Allocate blocks such that the only free spaces available are 3..6 and 9..10
337        // in order to prepare for the next test.
338        assert_eq!(alloc.allocate_range(3), Ok(0..3));
339        assert_eq!(alloc.allocate_range(3), Ok(3..6));
340        assert_eq!(alloc.allocate_range(3), Ok(6..9));
341        alloc.free_range(3..6);
342        assert_eq!(alloc.free_ranges, vec![3..6, 9..10]);
343        assert_eq!(
344            alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
345            vec![0..3, 6..9]
346        );
347        // Now request space that can be filled by 3..6 but should be filled by 9..10
348        // because 9..10 is a perfect fit.
349        assert_eq!(alloc.allocate_range(1), Ok(9..10));
350    }
351
352    #[test]
353    fn test_merge_neighbors() {
354        let mut alloc = RangeAllocator::new(0..9);
355        assert_eq!(alloc.allocate_range(3), Ok(0..3));
356        assert_eq!(alloc.allocate_range(3), Ok(3..6));
357        assert_eq!(alloc.allocate_range(3), Ok(6..9));
358        alloc.free_range(0..3);
359        alloc.free_range(6..9);
360        alloc.free_range(3..6);
361        assert_eq!(alloc.free_ranges, vec![0..9]);
362        assert!(alloc.allocated_ranges().eq(core::iter::empty()));
363    }
364}
range_alloc_arceos/lib.rs

range_alloc_arceos/
lib.rs
