buddy_system/
arena.rs

1use std::{iter::repeat_with, ops::Range, sync::mpsc, time::Instant};
2
3use generational_arena::{Arena, Index};
4
5use crate::{
6    buddy::{self, is_pow_of_two, Block, BlockState},
7    pretty_print::prettify,
8};
9
10/// NOT Copy or Clone, to make sure it's unique
11#[derive(Debug)]
12pub struct Allocation {
13    index: Index,
14    range: Range<usize>,
15    to_remove: mpsc::Sender<Index>,
16}
17
18impl Allocation {
19    pub fn range(&self) -> Range<usize> {
20        self.range.clone()
21    }
22}
23
24impl Drop for Allocation {
25    fn drop(&mut self) {
26        // don't panic here; leaking is not considered unsafe if the
27        // receiver doesn't get it for whatever reason
28        let _ = self.to_remove.send(self.index);
29    }
30}
31
32pub struct BuddyBookkeeping {
33    pub(crate) blocks: Arena<Block>,
34    pub(crate) root: Index,
35    to_remove_sender: mpsc::Sender<Index>,
36    to_remove_receiver: mpsc::Receiver<Index>,
37    min_block_size: usize,
38    max_block_size: usize,
39}
40
41impl BuddyBookkeeping {
42    pub fn new(size: usize, min_block_size: usize, max_block_size: usize) -> BuddyBookkeeping {
43        assert!(is_pow_of_two(size));
44        assert!(max_block_size <= size);
45        assert!(min_block_size <= max_block_size);
46
47        let mut new_arena = Arena::new();
48        let root = new_arena.insert(Block {
49            range: 0..size,
50            state: BlockState::Available,
51        });
52
53        let (sender, receiver) = mpsc::channel();
54
55        BuddyBookkeeping {
56            to_remove_sender: sender,
57            to_remove_receiver: receiver,
58            blocks: new_arena,
59            root: root,
60            min_block_size,
61            max_block_size,
62        }
63    }
64
65    pub fn alloc(&mut self, count: usize) -> Option<Allocation> {
66        let best_size = if 2_u32.pow(count.ilog2()) as usize == count {
67            count
68        } else {
69            2_u32.pow(count.ilog2() + 1) as usize
70        }
71        .max(self.min_block_size)
72        .min(self.max_block_size);
73
74        if best_size < count {
75            return None;
76        }
77
78        buddy::alloc(&mut self.blocks, self.root, best_size).map(|x| Allocation {
79            index: x,
80            range: (self.blocks[x].range.start)..(self.blocks[x].range.start + count),
81            to_remove: self.to_remove_sender.clone(),
82        })
83    }
84
85    pub fn tidy(&mut self) {
86        while let Ok(index) = self.to_remove_receiver.try_recv() {
87            buddy::dealloc(&mut self.blocks, index);
88        }
89
90        buddy::tidy(&mut self.blocks, self.root);
91    }
92
93    pub fn tidy_gas(&mut self, gas: usize) {
94        let mut gas = gas;
95
96        while let Ok(index) = self.to_remove_receiver.try_recv() {
97            if gas == 0 {
98                return;
99            }
100
101            buddy::dealloc(&mut self.blocks, index);
102
103            gas -= 1;
104        }
105
106        buddy::tidy_gas(&mut self.blocks, self.root, &mut gas);
107    }
108
109    pub fn tidy_timed(&mut self, deadline: Instant) {
110        while let Ok(index) = self.to_remove_receiver.try_recv() {
111            if Instant::now() >= deadline {
112                return;
113            }
114
115            buddy::dealloc(&mut self.blocks, index);
116        }
117
118        buddy::tidy_timed(&mut self.blocks, self.root, deadline);
119    }
120}
121
122pub struct BuddyArena<T> {
123    elements: Box<[T]>,
124    bookkeeping: BuddyBookkeeping,
125}
126
127impl<T> BuddyArena<T> {
128    pub fn new(size: usize, min_block_size: usize, max_block_size: usize) -> BuddyArena<T>
129    where
130        T: Default,
131    {
132        let elements_vec: Vec<T> = repeat_with(|| T::default()).take(size).collect();
133
134        BuddyArena {
135            elements: elements_vec.into(),
136            bookkeeping: BuddyBookkeeping::new(size, min_block_size, max_block_size),
137        }
138    }
139
140    pub fn bookkeeping(&self) -> &BuddyBookkeeping {
141        &self.bookkeeping
142    }
143
144    pub fn view(&self, a: &Allocation) -> &[T] {
145        &self.elements[a.range()]
146    }
147
148    pub fn view_mut(&mut self, a: &Allocation) -> &mut [T] {
149        &mut self.elements[a.range()]
150    }
151
152    pub fn alloc(&mut self, count: usize) -> Option<Allocation> {
153        self.bookkeeping.alloc(count)
154    }
155
156    pub fn tidy(&mut self) {
157        self.bookkeeping.tidy();
158    }
159
160    pub fn tidy_gas(&mut self, gas: usize) {
161        self.bookkeeping.tidy_gas(gas);
162    }
163
164    pub fn tidy_timed(&mut self, deadline: Instant) {
165        self.bookkeeping.tidy_timed(deadline);
166    }
167}
168
169#[test]
170fn test() {
171    let mut arena: BuddyArena<u8> = BuddyArena::new(2048, 8, 256);
172
173    let a1 = arena.alloc(64).unwrap();
174    let a2 = arena.alloc(24).unwrap();
175    let a3 = arena.alloc(2).unwrap();
176    let a4 = arena.alloc(7).unwrap();
177    let a5 = arena.alloc(31).unwrap();
178    let a6 = arena.alloc(60).unwrap();
179
180    println!("{:#?}", prettify(arena.bookkeeping()));
181
182    let string = "foobar";
183
184    let a = arena.alloc(string.bytes().len()).unwrap();
185    let view = arena.view_mut(&a);
186
187    view.copy_from_slice(string.as_bytes());
188
189    let str_view = std::str::from_utf8(view).unwrap();
190
191    dbg!(str_view);
192}
buddy_system/arena.rs

buddy_system/
arena.rs
