1#include <kassert.h>
2#include <mem/alloc.h>
3#include <string.h>
4#include <structures/id_space.h>
5
6static size_t id_space_get_data(struct rbt_node *node) {
7 return container_of(node, struct id_range, node)->start;
8}
9
10static int32_t id_space_cmp(const struct rbt_node *a,
11 const struct rbt_node *b) {
12 size_t l = id_space_get_data(node: (void *) a);
13 size_t r = id_space_get_data(node: (void *) b);
14 if (l < r) {
15 return -1;
16 }
17 if (l > r) {
18 return 1;
19 }
20 return 0;
21}
22
23struct id_space *id_space_init(uint64_t max_id) {
24 struct id_space *is = kmalloc(sizeof(*is), ALLOC_FLAGS_ZERO);
25 if (!is) {
26 return NULL;
27 }
28
29 rbt_init(t: &is->tree, get_data: id_space_get_data, compare: id_space_cmp);
30 spinlock_init(&is->lock);
31
32 is->reserve_free = NULL;
33 for (int i = 0; i < ID_RANGE_RESERVE_COUNT; i++) {
34 rbt_init_node(n: &is->reserve_pool[i].node);
35 is->reserve_pool[i].start = 0;
36 is->reserve_pool[i].length = 0;
37 is->reserve_pool[i].next = is->reserve_free;
38 is->reserve_free = &is->reserve_pool[i];
39 }
40
41 struct id_range *r = kmalloc(sizeof(*r), ALLOC_FLAGS_ZERO);
42 if (!r) {
43 if (is->reserve_free) {
44 r = is->reserve_free;
45 is->reserve_free = r->next;
46 } else {
47 return is;
48 }
49 }
50
51 r->start = 1;
52 r->length = max_id;
53 rbt_insert(tree: &is->tree, new_node: &r->node);
54
55 return is;
56}
57
58static struct id_range *id_range_alloc(struct id_space *is) {
59 SPINLOCK_ASSERT_HELD(&is->lock);
60 struct id_range *r = kmalloc(sizeof(*r), ALLOC_FLAGS_ZERO);
61 if (r) {
62 return r;
63 }
64
65 if (is->reserve_free) {
66 r = is->reserve_free;
67 is->reserve_free = r->next;
68 memset(r, 0, sizeof(*r));
69 return r;
70 }
71
72 return NULL;
73}
74
75static void id_range_free(struct id_space *is, struct id_range *r) {
76 SPINLOCK_ASSERT_HELD(&is->lock);
77 if ((uintptr_t) r >= (uintptr_t) &is->reserve_pool[0] &&
78 (uintptr_t) r < (uintptr_t) &is->reserve_pool[ID_RANGE_RESERVE_COUNT]) {
79 r->next = is->reserve_free;
80 is->reserve_free = r;
81 } else {
82 kfree(r);
83 }
84}
85
86uint64_t id_space_alloc(struct id_space *is) {
87 enum irql irql = spin_lock(&is->lock);
88
89 struct rbt_node *node = rbt_min(tree: &is->tree);
90 if (!node) {
91 spin_unlock(&is->lock, irql);
92 return 0;
93 }
94
95 struct id_range *range = rbt_entry(node, struct id_range, node);
96 uint64_t id = range->start;
97
98 if (range->length == 1) {
99 rbt_delete(tree: &is->tree, z: node);
100 id_range_free(is, r: range);
101 } else {
102 rbt_delete(tree: &is->tree, z: &range->node);
103 range->start++;
104 range->length--;
105 rbt_insert(tree: &is->tree, new_node: &range->node);
106 }
107
108 spin_unlock(&is->lock, irql);
109 return id;
110}
111
112void id_space_free(struct id_space *is, uint64_t id) {
113 enum irql irql = spin_lock(&is->lock);
114
115 struct rbt_node *node = is->tree.root;
116 struct id_range *prev = NULL;
117 struct id_range *next = NULL;
118
119 while (node) {
120 struct id_range *r = rbt_entry(node, struct id_range, node);
121 if (id < r->start) {
122 next = r;
123 node = node->left;
124 } else if (id > r->start + r->length - 1) {
125 prev = r;
126 node = node->right;
127 } else {
128 goto out;
129 }
130 }
131
132 bool merged_prev = false, merged_next = false;
133
134 if (prev && prev->start + prev->length == id) {
135 prev->length++;
136 merged_prev = true;
137 }
138
139 if (next && next->start == id + 1) {
140 if (merged_prev) {
141 prev->length += next->length;
142 rbt_delete(tree: &is->tree, z: &next->node);
143 id_range_free(is, r: next);
144 } else {
145 rbt_delete(tree: &is->tree, z: &next->node);
146 next->start = id;
147 next->length++;
148 rbt_insert(tree: &is->tree, new_node: &next->node);
149 }
150 merged_next = true;
151 }
152
153 if (!merged_prev && !merged_next) {
154 struct id_range *new_range = id_range_alloc(is);
155 if (!new_range) {
156 goto out;
157 }
158
159 new_range->start = id;
160 new_range->length = 1;
161 rbt_insert(tree: &is->tree, new_node: &new_range->node);
162 }
163
164out:
165 spin_unlock(&is->lock, irql);
166}
167
168uint64_t id_space_alloc_range(struct id_space *is, uint64_t count) {
169 if (count == 0) {
170 return 0;
171 }
172 if (count == 1) {
173 return id_space_alloc(is);
174 }
175
176 enum irql irql = spin_lock(&is->lock);
177
178 struct rbt_node *node;
179 struct id_range *range = NULL;
180
181 rbt_for_each(node, &is->tree) {
182 struct id_range *r = rbt_entry(node, struct id_range, node);
183 if (r->length >= count) {
184 range = r;
185 break;
186 }
187 }
188
189 if (!range) {
190 spin_unlock(&is->lock, irql);
191 return 0;
192 }
193
194 uint64_t id = range->start;
195
196 if (range->length == count) {
197 rbt_delete(tree: &is->tree, z: &range->node);
198 id_range_free(is, r: range);
199 } else {
200 rbt_delete(tree: &is->tree, z: &range->node);
201 range->start += count;
202 range->length -= count;
203 rbt_insert(tree: &is->tree, new_node: &range->node);
204 }
205
206 spin_unlock(&is->lock, irql);
207 return id;
208}
209
210void id_space_free_range(struct id_space *is, uint64_t start, uint64_t count) {
211 for (uint64_t i = 0; i < count; i++) {
212 id_space_free(is, id: start + i);
213 }
214}
215
216void id_space_destroy(struct id_space *is) {
217 if (!is) {
218 return;
219 }
220
221 enum irql irql = spin_lock(&is->lock);
222 struct rbt_node *node;
223 struct rbt_node *tmp;
224
225 rbt_for_each_safe(node, tmp, &is->tree) {
226 struct id_range *r = rbt_entry(node, struct id_range, node);
227 rbt_delete(tree: &is->tree, z: node);
228 id_range_free(is, r);
229 }
230 spin_unlock(&is->lock, irql);
231
232 kfree(is);
233}
234