From ad8c66bb920b6fcccf86eba9a0ae642889a4656d Mon Sep 17 00:00:00 2001
From: Matthew Sotoudeh <matthew@masot.net>
Date: Fri, 17 May 2024 13:39:02 -0700
Subject: simplify the reservation logic

---
 magic_buddy/magic_buddy.c | 79 +++++++++++++++++++----------------------------
 1 file changed, 31 insertions(+), 48 deletions(-)

diff --git a/magic_buddy/magic_buddy.c b/magic_buddy/magic_buddy.c
index f3dac00..06f5735 100644
--- a/magic_buddy/magic_buddy.c
+++ b/magic_buddy/magic_buddy.c
@@ -142,58 +142,41 @@ static struct free_block *rhs_child_of(struct free_block *block, size_t logsize,
 // take great pains to avoid writing to the region that is being reserved
 // (e.g., in case it is device MMIO).
 int reserve(void *start, size_t size, struct buddy *state) {
-    assert(size);
-    void *end = (void*)((uint8_t*)start + size);
-
-    // (1) iterate down to find the first free node to the left of @start
-    struct free_block *parent = (void*)state->base;
-    size_t parent_logsize = state->root_logsize;
-
-    while (!isfree(parent, parent_logsize, state)) {
-        // pick the proper child and recurse
-        struct free_block *rhs_child
-            = rhs_child_of(parent, parent_logsize, state);
-        if ((void*)rhs_child <= start) parent = rhs_child;
-        // TODO: make sure they don't try to reserve below the min size here?
-        if (!parent_logsize--) return 0;
+    // (1) find the first free block to the left of start
+    uint8_t *base = state->base;
+    size_t virtual_start = (uint8_t*)start - base;
+    struct free_block *block = (void*)(base + virtual_start);
+    // repeatedly zero out least significant bits until we find a free block
+    while (memcmp(block->magic, state->magic, sizeof(state->magic))) {
+        if (!virtual_start) return 0;
+        virtual_start = virtual_start & (virtual_start - 1);
+        block = (void*)(base + virtual_start);
     }
-    // parent is free!
-    void *parent_end = (void*)((uint8_t*)parent + (1 << parent_logsize));
-    // if the size has no chance of fitting, return 0.
-    if (parent_end < end) return 0;
-    // otherwise, split this parent
-    pop(parent);
-    // TODO: any way to avoid this write? wouldn't be necessary if we're on the
-    // rhs of the parent ...
-    memset(parent, 0, sizeof(struct free_block));
-    while (1) {
-        // check whether the child could fit it.
-        struct free_block *rhs_child
-            = rhs_child_of(parent, parent_logsize, state);
-        struct free_block *next_child
-            = ((void*)rhs_child <= start) ? rhs_child : parent;
-        void *child_end
-            = (void*)((uint8_t*)next_child + (1 << (parent_logsize - 1)));
-
-        // if the child *cannot* fit it, stop by leaving this parent allocated.
-        if (child_end < end) return 1;
-
-        // if the child *can* fit it, split this parent into two children.
-        struct free_block *lhs_child = parent;
-        if ((void*)rhs_child <= start) {
-            // the lhs child will be free
-            makefree(lhs_child, state);
-            push(lhs_child, parent_logsize - 1, state);
-            parent = rhs_child;
-        } else {
-            // the rhs child will be free
+
+    // (2) check whether the block fits it
+    void *end = (void*)((uint8_t*)start + size);
+    void *block_end = ((uint8_t*)block + (1 << block->logsize));
+    if (block_end < end) return 0;
+
+    // (3) split the block until we get a tight fit
+    int needs_zero = 1;
+    size_t min_logsize = size2log(sizeof(struct free_block), 1);
+    pop(block);
+    for (size_t logsize = block->logsize; logsize > min_logsize;) {
+        struct free_block *rhs_child = rhs_child_of(block, logsize, state);
+        if ((void*)rhs_child <= start) {        // move right
+            block = rhs_child;
+            makefree(block, state);
+            push(block, --logsize, state);
+            needs_zero = 0;
+        } else {                                // move left
+            if (end > (void*)rhs_child) break;
             makefree(rhs_child, state);
-            push(rhs_child, parent_logsize - 1, state);
-            parent = lhs_child;
+            push(rhs_child, --logsize, state);
         }
-        parent_logsize--;
     }
-    assert(!"unreachable");
+    if (needs_zero) memset(block, 0, sizeof(struct free_block));
+    return 1;
 }
 
 static void *_naive_reallocate(void *old, size_t old_size, size_t new_size,
-- 
cgit v1.2.3