rustbuddy/
lib.rs

1use std::vec;
2
3const NODE_UNUSED: u8 = 0;
4const NODE_USED: u8 = 1;
5const NODE_SPLIT: u8 = 2;
6const NODE_FULL: u8 = 3;
7
8
9pub struct BuddyAllocator {
10    levels: usize,
11    pub tree: vec::Vec<u8>,
12}
13
14impl BuddyAllocator {
15    pub fn new(levels: usize) -> BuddyAllocator {
16        let size: usize = (1 << levels + 1) - 1;
17        return BuddyAllocator{
18            levels: levels,
19            tree: vec![NODE_UNUSED; size],
20        };
21    }
22    
23    // Takes a size (# of blocks requested) and returns an index offset
24    pub fn allocate(&mut self, s: usize) -> isize {
25        // Get the number of blocks requested
26        let requested_blocks: f64;
27        if s == 0 {
28            requested_blocks = 1.0;
29        } else {
30            requested_blocks = s.next_power_of_two() as f64;
31        }
32        let requested_level = requested_blocks.log(2.0) as usize;
33        if requested_level > self.levels {
34            return -1;
35        }
36
37        // start at index 0 and move in
38        let mut index = 0;
39        let mut current_level = self.levels;
40        'forward: loop {
41            let has_buddy = index & 1 == 1;
42            if current_level != requested_level {
43                match self.tree[index] { 
44                    NODE_USED | NODE_FULL => {
45                        // Check the buddy node if we haven't already
46                        if has_buddy {
47                            index += 1;
48                        }
49                        continue 'forward;
50                    }
51                    NODE_UNUSED => {
52                        // Split the node and descend
53                        self.tree[index] = NODE_SPLIT;
54                        index = index * 2 + 1;
55                        current_level -= 1;
56                        continue 'forward;
57                    }
58                    NODE_SPLIT => {
59                        // Just descend
60                        index = index * 2 + 1;
61                        current_level -= 1;
62                        continue 'forward;
63                    }
64                    _ => panic!("unknkown type {}", self.tree[index])
65                }
66            } else {
67                // Requested level and current level match up
68                if self.tree[index] == NODE_UNUSED {
69                    self.tree[index] = NODE_USED;
70                    // Recursively check if parents are full and mark them as such
71                    self.update_parents((index + 1) / 2 - 1);
72                    break 'forward;
73                }
74            }
75            // Check buddy node if we haven't already
76            if has_buddy {
77                index += 1;
78                continue 'forward;
79            }
80            // Backtrack if we reach a level match AND we've checked both nodes
81            'backward: loop {
82                index = (index + 1) / 2 - 1;
83                current_level += 1;
84                let has_buddy_inner = index & 1 == 1;
85                if has_buddy_inner {
86                    index += 1;
87                    break 'backward;
88                }
89            }
90        }
91
92        return index as isize;
93    }
94
95    pub fn free(&mut self, index_offset: usize) {
96        if index_offset > self.tree.len() - 1 {
97            panic!("offset {} is > length of tree {}", index_offset, self.tree.len());
98        }
99        // Recursively free and combine nodes
100        self.free_and_combine(index_offset);
101        
102        // Recursively update parents
103        self.update_parents((index_offset + 1) / 2 - 1);
104    }
105
106    fn free_and_combine(&mut self, index: usize) {
107        self.tree[index] = NODE_UNUSED;
108        // We are already at the top of the tree, we're done
109        if index == 0 {
110            return;
111        }
112        let other_node: usize;
113        let has_right_buddy = (index & 1) == 1;
114        if has_right_buddy {
115            other_node = index + 1;
116        } else {
117            other_node = index - 1;
118        }
119        // Recursively combine nodes
120        if self.tree[other_node] == NODE_UNUSED {
121            self.free_and_combine((index + 1) / 2 - 1);
122        }
123        return;
124    }
125
126    // Propagate changes up to parent nodes
127    fn update_parents(&mut self, index: usize) {
128        // Check both child nodes to see if they are both either FULL or USED
129        let left_child = index * 2 + 1;
130        let right_child = index * 2 + 2;
131        let left_child_used_or_full = self.tree[left_child] == NODE_FULL || self.tree[left_child] == NODE_USED;
132        let right_child_used_or_full = self.tree[right_child] == NODE_FULL || self.tree[right_child] == NODE_USED;
133        if left_child_used_or_full && right_child_used_or_full {
134            // Both children USED or FULL
135            self.tree[index] = NODE_FULL;
136        } else if self.tree[left_child] == NODE_UNUSED && self.tree[right_child] == NODE_UNUSED {
137            // Both children are UNUSED
138            self.tree[index] = NODE_UNUSED;
139        } else {
140            // Default to split node if neither FULL or UNUSED
141            self.tree[index] = NODE_SPLIT;
142        }
143        // We're at the top of the tree, we're done
144        if index == 0 {
145            return;
146        }
147        self.update_parents((index + 1) / 2 - 1);
148    }
149
150    pub fn dump(&self) -> String { 
151        let mut out = "".to_string();
152        let mut row = "".to_string();
153        let mut level = 0;
154        let mut index = 0;
155        loop {
156            if index == self.tree.len() {
157                break
158            }
159            match self.tree[index] {
160                NODE_USED => row += "U",
161                NODE_UNUSED => row += "O",
162                NODE_SPLIT => row += "S",
163                NODE_FULL => row += "F",
164                _ => panic!("unknown node type {}", self.tree[index]),
165            }
166            if row.len() == 1 << level {
167                out += &(row + "\n");
168                row = "".to_string();
169                level += 1;
170            }
171            index += 1;
172        }
173        return out;
174    }
175}
rustbuddy/lib.rs

rustbuddy/
lib.rs
