buddy_system/
buddy.rs

1//! Buddy system implementations.
2//!
3//! Note: everything in this module is unchecked. It shouldn't panic (the only `unwraps`
4//! should be unreachable unless an internal invariant is broken), but it won't behave as
5//! expected if it's given the wrong inputs or state.
6
7use std::{cmp::Ordering, ops::Range, time::Instant};
8
9use generational_arena::{Arena, Index};
10
11pub(crate) fn is_pow_of_two(x: usize) -> bool {
12    (x != 0) && ((x & (x - 1)) == 0)
13}
14
15pub(crate) enum BlockState {
16    Split(Index, Index),
17    Available,
18    Occupied,
19}
20
21pub struct Block {
22    pub(crate) range: Range<usize>,
23    pub(crate) state: BlockState,
24}
25
26/// Assumes `desired_size` is a power of 2
27pub fn alloc(arena: &mut Arena<Block>, block_index: Index, desired_size: usize) -> Option<Index> {
28    debug_assert!(is_pow_of_two(desired_size));
29
30    let block = &arena[block_index];
31
32    match block.range.len().cmp(&desired_size) {
33        Ordering::Less => None,
34        Ordering::Equal => {
35            if let BlockState::Available = block.state {
36                arena[block_index].state = BlockState::Occupied;
37
38                Some(block_index)
39            } else {
40                None
41            }
42        }
43        Ordering::Greater => match block.state {
44            BlockState::Occupied => None,
45            BlockState::Available => {
46                let first_range = (block.range.start)..(block.range.start + block.range.len() / 2);
47                let second_range = (block.range.start + block.range.len() / 2)..(block.range.end);
48
49                let first = arena.insert(Block {
50                    range: first_range,
51                    state: BlockState::Available,
52                });
53
54                let second = arena.insert(Block {
55                    range: second_range,
56                    state: BlockState::Available,
57                });
58
59                arena[block_index].state = BlockState::Split(first, second);
60
61                alloc(arena, first, desired_size)
62            }
63            BlockState::Split(first_index, second_index) => {
64                if let Some(result) = alloc(arena, first_index, desired_size) {
65                    Some(result)
66                } else if let Some(result) = alloc(arena, second_index, desired_size) {
67                    Some(result)
68                } else {
69                    None
70                }
71            }
72        },
73    }
74}
75
76pub(crate) fn dealloc(arena: &mut Arena<Block>, block_index: Index) {
77    arena[block_index].state = BlockState::Available;
78}
79
80#[repr(transparent)]
81pub struct IsAvailable(bool);
82
83pub fn tidy(arena: &mut Arena<Block>, block_index: Index) -> IsAvailable {
84    // go through and merge
85    let block = &arena[block_index];
86
87    match block.state {
88        BlockState::Split(first, second) => {
89            let first_available = tidy(arena, first).0;
90            let second_available = tidy(arena, second).0;
91
92            if first_available && second_available {
93                arena.remove(first).unwrap();
94                arena.remove(second).unwrap();
95
96                arena[block_index].state = BlockState::Available;
97
98                IsAvailable(true)
99            } else {
100                IsAvailable(false)
101            }
102        }
103        BlockState::Available => IsAvailable(true),
104        BlockState::Occupied => IsAvailable(false),
105    }
106}
107
108pub fn tidy_gas(arena: &mut Arena<Block>, block_index: Index, gas: &mut usize) -> IsAvailable {
109    if *gas == 0 {
110        return IsAvailable(false);
111    }
112
113    *gas -= 1;
114
115    // go through and merge
116    let block = &arena[block_index];
117
118    match block.state {
119        BlockState::Split(first, second) => {
120            let first_available = tidy_gas(arena, first, gas).0;
121            let second_available = tidy_gas(arena, second, gas).0;
122
123            if first_available && second_available {
124                arena.remove(first).unwrap();
125                arena.remove(second).unwrap();
126
127                arena[block_index].state = BlockState::Available;
128
129                IsAvailable(true)
130            } else {
131                IsAvailable(false)
132            }
133        }
134        BlockState::Available => IsAvailable(true),
135        BlockState::Occupied => IsAvailable(false),
136    }
137}
138
139pub fn tidy_timed(arena: &mut Arena<Block>, block_index: Index, deadline: Instant) -> IsAvailable {
140    if Instant::now() >= deadline {
141        // return not available so the recursion chain stops
142        return IsAvailable(false);
143    }
144
145    // go through and merge
146    let block = &arena[block_index];
147
148    match block.state {
149        BlockState::Split(first, second) => {
150            let first_available = tidy_timed(arena, first, deadline).0;
151            let second_available = tidy_timed(arena, second, deadline).0;
152
153            if first_available && second_available {
154                arena.remove(first).unwrap();
155                arena.remove(second).unwrap();
156
157                arena[block_index].state = BlockState::Available;
158
159                IsAvailable(true)
160            } else {
161                IsAvailable(false)
162            }
163        }
164        BlockState::Available => IsAvailable(true),
165        BlockState::Occupied => IsAvailable(false),
166    }
167}
buddy_system/buddy.rs

buddy_system/
buddy.rs
