1#include <console/panic.h>
2#include <kassert.h>
3#include <math/align.h>
4#include <math/div.h>
5#include <mem/address_range.h>
6#include <mem/alloc.h>
7#include <mem/alloc_or_die.h>
8#include <mem/hhdm.h>
9#include <mem/page.h>
10#include <mem/pmm.h>
11#include <mem/vas.h>
12#include <smp/core.h>
13#include <string.h>
14
15static size_t vas_cpu_id() {
16 if (global.current_bootstage < BOOTSTAGE_MID_MP)
17 return 0;
18
19 return smp_core_id();
20}
21
22static size_t vas_range_get_data(struct rbt_node *n) {
23 return container_of(n, struct vas_range, node)->start;
24}
25
26static int32_t vas_range_cmp(const struct rbt_node *a,
27 const struct rbt_node *b) {
28 vaddr_t l = vas_range_get_data(n: (void *) a);
29 vaddr_t r = vas_range_get_data(n: (void *) b);
30 return (l > r) - (l < r);
31}
32
33static struct vas_range *vasrange_alloc(struct vas_local_tree *lt) {
34 return fixed_size_alloc(fsr: &lt->fsr);
35}
36
37static void vasrange_free(struct vas_local_tree *lt, struct vas_range *r) {
38 return fixed_size_free(fsr: &lt->fsr, obj: r);
39}
40
41void vas_reclaim_freelist_pages(struct vas_local_tree *lt) {
42 fixed_size_reclaim_freelist_pages(fsr: &lt->fsr);
43}
44
45static void vas_range_fsr_init(void *n) {
46 struct vas_range *r = n;
47 rbt_init_node(n: &r->node);
48 r->length = 0;
49 r->start = 0;
50}
51
52static void vas_local_tree_init(struct vas_local_tree *lt) {
53 spinlock_init(lock: &lt->lock);
54 rbt_init(t: &lt->tree, get_data: vas_range_get_data, compare: vas_range_cmp);
55 struct fixed_size_range_attributes attrs = {
56 .init_obj = vas_range_fsr_init,
57 .deinit_obj = NULL,
58 .bootstrap_mode = true,
59 .obj_align = _Alignof(struct vas_range),
60 .obj_size = sizeof(struct vas_range),
61 };
62
63 fixed_size_range_init(fsr: &lt->fsr, attrs: &attrs);
64}
65
66static vaddr_t tree_alloc(struct vas_local_tree *lt, size_t size,
67 size_t align) {
68 struct rbt_node *node = rbt_min(tree: &lt->tree);
69
70 while (node) {
71 struct vas_range *gap = rbt_entry(node, struct vas_range, node);
72
73 vaddr_t aligned = ALIGN_UP(gap->start, align);
74 vaddr_t end = aligned + size;
75
76 if (end <= gap->start + gap->length) {
77 rbt_delete(tree: &lt->tree, z: &gap->node);
78
79 if (aligned > gap->start) {
80 struct vas_range *left = vasrange_alloc(lt);
81 if (!left)
82 goto fail;
83 left->start = gap->start;
84 left->length = aligned - gap->start;
85 rbt_insert(tree: &lt->tree, new_node: &left->node);
86 }
87
88 if (end < gap->start + gap->length) {
89 struct vas_range *right = vasrange_alloc(lt);
90 if (!right)
91 goto fail;
92 right->start = end;
93 right->length = (gap->start + gap->length) - end;
94 rbt_insert(tree: &lt->tree, new_node: &right->node);
95 }
96
97 vasrange_free(lt, r: gap);
98 lt->total_free -= size;
99 return aligned;
100
101 fail:
102 rbt_insert(tree: &lt->tree, new_node: &gap->node);
103 return 0;
104 }
105
106 node = rbt_next(node);
107 }
108
109 return 0;
110}
111
112static void tree_free(struct vas_local_tree *lt, vaddr_t addr, size_t size) {
113 vaddr_t start = addr;
114 vaddr_t end = addr + size;
115
116 struct rbt_node *node = lt->tree.root;
117 struct vas_range *prev = NULL;
118 struct vas_range *next = NULL;
119
120 while (node) {
121 struct vas_range *g = rbt_entry(node, struct vas_range, node);
122
123 if (end <= g->start) {
124 next = g;
125 node = node->left;
126 } else if (start >= g->start + g->length) {
127 prev = g;
128 node = node->right;
129 } else {
130 panic("vas_free: overlap/double free at %p+%zu", (void *) addr,
131 size);
132 }
133 }
134
135 struct vas_range *reuse = NULL;
136
137 if (prev && prev->start + prev->length == start) {
138 start = prev->start;
139 size += prev->length;
140 rbt_delete(tree: &lt->tree, z: &prev->node);
141 reuse = prev;
142 }
143
144 if (next && end == next->start) {
145 size += next->length;
146 rbt_delete(tree: &lt->tree, z: &next->node);
147 if (reuse)
148 vasrange_free(lt, r: next);
149 else
150 reuse = next;
151 }
152
153 if (!reuse)
154 reuse = vasrange_alloc(lt);
155
156 reuse->start = start;
157 reuse->length = size;
158 rbt_insert(tree: &lt->tree, new_node: &reuse->node);
159
160 lt->total_free += (end - addr);
161}
162
163static bool vas_pull_chunk(struct vas *vas, struct vas_local_tree *lt) {
164 struct vas_local_tree *gl = &vas->global;
165 size_t chunk = vas->chunk_size;
166
167 enum irql irql = spin_lock(lock: &gl->lock);
168
169 /* Try to allocate a full chunk from the global tree */
170 vaddr_t base = tree_alloc(lt: gl, size: chunk, align: chunk);
171
172 if (!base) {
173 struct rbt_node *node = rbt_min(tree: &gl->tree);
174 struct vas_range *best = NULL;
175
176 while (node) {
177 struct vas_range *g = rbt_entry(node, struct vas_range, node);
178 if (!best || g->length > best->length)
179 best = g;
180 node = rbt_next(node);
181 }
182
183 if (best && best->length >= PAGE_SIZE) {
184 /* Take the whole gap */
185 size_t take = best->length;
186 base = best->start;
187
188 rbt_delete(tree: &gl->tree, z: &best->node);
189 gl->total_free -= take;
190 vasrange_free(lt: gl, r: best);
191
192 spin_unlock(lock: &gl->lock, old: irql);
193
194 /* Insert into local tree */
195 enum irql lirql = spin_lock(lock: &lt->lock);
196
197 struct vas_range *nr = vasrange_alloc(lt);
198 if (!nr) {
199 spin_unlock(lock: &lt->lock, old: lirql);
200 /* Return the chunk back to global */
201 irql = spin_lock(lock: &gl->lock);
202 tree_free(lt: gl, addr: base, size: take);
203 spin_unlock(lock: &gl->lock, old: irql);
204 return false;
205 }
206
207 nr->start = base;
208 nr->length = take;
209 rbt_insert(tree: &lt->tree, new_node: &nr->node);
210 lt->total_free += take;
211
212 spin_unlock(lock: &lt->lock, old: lirql);
213 return true;
214 }
215
216 spin_unlock(lock: &gl->lock, old: irql);
217 return false;
218 }
219
220 spin_unlock(lock: &gl->lock, old: irql);
221
222 /* Insert the chunk as a single gap into the local tree */
223 enum irql lirql = spin_lock(lock: &lt->lock);
224
225 struct vas_range *nr = vasrange_alloc(lt);
226 if (!nr) {
227 spin_unlock(lock: &lt->lock, old: lirql);
228 /* Return chunk to global */
229 irql = spin_lock(lock: &gl->lock);
230 tree_free(lt: gl, addr: base, size: chunk);
231 spin_unlock(lock: &gl->lock, old: irql);
232 return false;
233 }
234
235 nr->start = base;
236 nr->length = chunk;
237 rbt_insert(tree: &lt->tree, new_node: &nr->node);
238 lt->total_free += chunk;
239
240 spin_unlock(lock: &lt->lock, old: lirql);
241 return true;
242}
243
244/* faster approach: maintain a sorted interval list per-CPU of chunk
245 * boundaries and binary-search it, but the simple version works first */
246static ssize_t vas_find_owner(struct vas *vas, vaddr_t addr) {
247 size_t i;
248 for_each_cpu_id(i) {
249 struct vas_local_tree *lt = &vas->local[i];
250 enum irql irql = spin_lock(lock: &lt->lock);
251
252 struct rbt_node *node = lt->tree.root;
253 bool found = false;
254
255 while (node) {
256 struct vas_range *g = rbt_entry(node, struct vas_range, node);
257
258 if (addr < g->start) {
259 /* addr could be just below this gap */
260 if (addr + PAGE_SIZE >= g->start) {
261 /* this tree likely owns it */
262 found = true;
263 break;
264 }
265 node = node->left;
266 } else if (addr >= g->start + g->length) {
267 /* addr is above this gap */
268 if (addr < g->start + g->length + PAGE_SIZE) {
269 found = true;
270 break;
271 }
272 node = node->right;
273 } else {
274 /* double free */
275 found = true;
276 break;
277 }
278 }
279
280 spin_unlock(lock: &lt->lock, old: irql);
281
282 if (found)
283 return (ssize_t) i;
284 }
285
286 return -1;
287}
288
289vaddr_t vas_alloc(struct vas *vas, size_t size, size_t align) {
290 uint32_t cpu = vas_cpu_id();
291 struct vas_local_tree *lt = &vas->local[cpu];
292 vaddr_t result;
293
294 enum irql irql = spin_lock(lock: &lt->lock);
295 result = tree_alloc(lt, size, align);
296 spin_unlock(lock: &lt->lock, old: irql);
297
298 if (result)
299 return result;
300
301 if (!vas_pull_chunk(vas, lt))
302 return 0;
303
304 irql = spin_lock(lock: &lt->lock);
305 result = tree_alloc(lt, size, align);
306 spin_unlock(lock: &lt->lock, old: irql);
307
308 return result;
309}
310
311void vas_free(struct vas *vas, vaddr_t addr, size_t size) {
312
313 ssize_t owner = vas_find_owner(vas, addr);
314 struct vas_local_tree *lt;
315
316 if (owner >= 0) {
317 lt = &vas->local[owner];
318 } else {
319 lt = &vas->local[vas_cpu_id()];
320 }
321
322 enum irql irql = spin_lock(lock: &lt->lock);
323 tree_free(lt, addr, size);
324 spin_unlock(lock: &lt->lock, old: irql);
325}
326
327/* True if addr lands inside a free gap held by this one tree. The tree is keyed
328 * by gap start, so this is a plain BST containment search (the same shape as
329 * the inner loops of tree_free / vas_find_owner). Caller holds lt->lock. */
330static bool tree_addr_is_free(struct vas_local_tree *lt, vaddr_t addr) {
331 struct rbt_node *node = lt->tree.root;
332 while (node) {
333 struct vas_range *g = rbt_entry(node, struct vas_range, node);
334
335 if (addr < g->start)
336 node = node->left;
337 else if (addr >= g->start + g->length)
338 node = node->right;
339 else
340 return true; /* g->start <= addr < g->start + g->length */
341 }
342 return false;
343}
344
345bool vas_vaddr_is_allocated(struct vas *vas, vaddr_t addr) {
346 if (!vas_vaddr_in_vas(vas, vaddr: addr))
347 return false;
348
349 enum irql irql = spin_lock(lock: &vas->global.lock);
350 bool free_here = tree_addr_is_free(lt: &vas->global, addr);
351 spin_unlock(lock: &vas->global.lock, old: irql);
352 if (free_here)
353 return false;
354
355 size_t i;
356 for_each_cpu_id(i) {
357 struct vas_local_tree *lt = &vas->local[i];
358 irql = spin_lock(lock: &lt->lock);
359 free_here = tree_addr_is_free(lt, addr);
360 spin_unlock(lock: &lt->lock, old: irql);
361 if (free_here)
362 return false;
363 }
364
365 return true;
366}
367
368static void vas_space_init_common(struct vas *vas, vaddr_t base,
369 vaddr_t limit) {
370 vas->base = base;
371 vas->limit = limit;
372 vas->chunk_size = VAS_CHUNK_SIZE;
373
374 /* Init global tree */
375 vas_local_tree_init(lt: &vas->global);
376
377 enum irql irql = spin_lock(lock: &vas->global.lock);
378
379 struct vas_range *g = alloc_or_die(vasrange_alloc(&vas->global));
380
381 g->start = base;
382 g->length = limit - base;
383 rbt_insert(tree: &vas->global.tree, new_node: &g->node);
384 vas->global.total_free = limit - base;
385
386 spin_unlock(lock: &vas->global.lock, old: irql);
387
388 for (uint32_t i = 0; i < global.core_count; i++)
389 vas_local_tree_init(lt: &vas->local[i]);
390}
391
392struct vas *vas_bootstrap(vaddr_t base, vaddr_t limit) {
393 if (global.current_bootstage >= BOOTSTAGE_LATE)
394 log_warn_once("vas_space_bootstrap called after bootstage %s",
395 bootstage_str[BOOTSTAGE_LATE]);
396
397 uint32_t ncpus = global.core_count;
398 size_t total_size =
399 sizeof(struct vas) + ncpus * sizeof(struct vas_local_tree);
400 size_t pages_needed = DIV_ROUND_UP(total_size, PAGE_SIZE);
401
402 uintptr_t phys = alloc_or_die(pmm_alloc_page(pages_needed));
403 uintptr_t virt = hhdm_paddr_to_vaddr(p: phys);
404 memset((void *) virt, 0, PAGE_SIZE);
405
406 struct vas *vas = (struct vas *) virt;
407 vas->local =
408 (struct vas_local_tree *) ((uint8_t *) vas + sizeof(struct vas));
409
410 vas_space_init_common(vas, base, limit);
411 return vas;
412}
413
414struct vas *vas_create(vaddr_t base, vaddr_t limit) {
415 uint32_t ncpus = global.core_count;
416 size_t total_size =
417 sizeof(struct vas) + ncpus * sizeof(struct vas_local_tree);
418
419 struct vas *vas =
420 kmalloc(total_size, ALLOC_FLAGS_ZERO,
421 ALLOC_BEHAVIOR_NORMAL | ALLOC_BEHAVIOR_FLAG_MINIMAL);
422 if (!vas)
423 return NULL;
424
425 vas->local =
426 (struct vas_local_tree *) ((uint8_t *) vas + sizeof(struct vas));
427
428 vas_space_init_common(vas, base, limit);
429 return vas;
430}
431
432void *vas_map(struct vas *vas, paddr_t paddr, size_t len, uint64_t flags,
433 enum vmm_flags vflags) {
434 size_t pages = DIV_ROUND_UP(len, PAGE_SIZE);
435 vaddr_t vaddr = vas_alloc(vas, PAGE_SIZE * pages, PAGE_SIZE);
436 if (!vaddr)
437 return NULL;
438
439 void *ret = vmm_map(paddr, vaddr, len, flags, vflags);
440 if (!ret)
441 vas_free(vas, addr: vaddr, PAGE_SIZE * pages);
442
443 return ret;
444}
445
446void vas_unmap(struct vas *vas, void *vaddr, size_t len) {
447 size_t pages = DIV_ROUND_UP(len, PAGE_SIZE);
448 vmm_unmap(addr: vaddr, PAGE_SIZE * pages, vflags: VMM_FLAG_NONE);
449 vas_free(vas, addr: (vaddr_t) vaddr, PAGE_SIZE * pages);
450}
451
452struct vas *vas_from(struct address_range *ar) {
453 return vas_create(base: ar->base, limit: ar->base + ar->size);
454}
455
456struct vas *vas_bootstrap_from(struct address_range *ar) {
457 return vas_bootstrap(base: ar->base, limit: ar->base + ar->size);
458}
459
460static void vas_dump_local_tree(struct vas_local_tree *lt, bool take_lock) {
461 enum irql irql = 0;
462 if (take_lock)
463 irql = spin_lock(lock: &lt->lock);
464
465 size_t count = 0;
466 size_t largest = 0;
467 vaddr_t largest_start = 0;
468
469 for (struct rbt_node *n = rbt_min(tree: &lt->tree); n; n = rbt_next(node: n)) {
470 struct vas_range *r = rbt_entry(n, struct vas_range, node);
471 printf(format: " [%zu] %p .. %p len=%zu (%zu KiB)\n", count,
472 (void *) r->start, (void *) (r->start + r->length), r->length,
473 r->length >> 10);
474 if (r->length > largest) {
475 largest = r->length;
476 largest_start = r->start;
477 }
478 count++;
479 }
480
481 printf(
482 format: " %zu free range(s), total_free=%zu (%zu KiB), largest=%zu @ %p\n",
483 count, lt->total_free, lt->total_free >> 10, largest,
484 (void *) largest_start);
485
486 if (take_lock)
487 spin_unlock(lock: &lt->lock, old: irql);
488}
489
490void vas_space_dump(struct vas *vas) {
491 size_t span = (size_t) (vas->limit - vas->base);
492 printf(format: "vas_space %p: base=%p limit=%p span=%zu (%zu MiB) chunk_size=%zu\n",
493 (void *) vas, (void *) vas->base, (void *) vas->limit, span,
494 span >> 20, vas->chunk_size);
495
496 printf(format: " global tree:\n");
497 vas_dump_local_tree(lt: &vas->global, true);
498
499 for (uint32_t i = 0; i < global.core_count; i++) {
500 printf(format: " cpu[%u] tree:\n", i);
501 vas_dump_local_tree(lt: &vas->local[i], true);
502 }
503}
504