From 35674243f07216ce94c2680fde13ee9df48d766c Mon Sep 17 00:00:00 2001
From: Matthew Sotoudeh <matthew@masot.net>
Date: Thu, 16 May 2024 10:39:47 -0700
Subject: support for moving the allocation pool

---
 README        | 22 +++-------------------
 buddy.h       |  5 +++++
 magic_buddy.c | 60 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 19 deletions(-)

diff --git a/README b/README
index debd59b..d85de3f 100644
--- a/README
+++ b/README
@@ -54,27 +54,11 @@ in the computer's address size.
 
 ### Usage
 
+Usage example in main.c.
+
 First, initialize the library by giving it both the base + size for the
 allocation pool to be managed, along with some "very random" (or otherwise
-secret) data. The following should work:
-
-    // First, get some secret random data.
-    uint8_t secret[MAGIC_COOKIE_BYTES];
-    int fd = open("/dev/urandom", O_RDONLY);
-    assert(count == read(fd, secret, count));
-    close(fd);
-
-    // Then initialize the buddy allocator.
-    init_buddy(malloc(1026 * 1024), 1024 * 1024, secret);
-
-Now, you can allocate and liberate as normal:
-
-    // Get a region of 128 bytes
-    uint8_t *x = allocate(128);
-    // Write into it
-    x[123] = 1234;
-    // Liberate it
-    liberate(x, 128);
+secret) data. Then, you can allocate and liberate as normal.
 
 ### License
 
diff --git a/buddy.h b/buddy.h
index 0b4de09..ad4f4aa 100644
--- a/buddy.h
+++ b/buddy.h
@@ -24,3 +24,8 @@ void debug_buddy(struct buddy *state);
 void *reallocate(void *old, size_t old_size, size_t new_size,
                  struct buddy *state);
 int reserve(void *start, size_t size, struct buddy *state);
+
+// grow can also be used to move the buddy. grow can never fail.
+void grow_buddy(uint8_t *new_base, size_t new_size, struct buddy *state);
+// 0 -> failure. always in-place. to do out-of-place, first shrink then grow.
+int shrink_buddy(size_t new_size, struct buddy *state);
diff --git a/magic_buddy.c b/magic_buddy.c
index f87ae35..a741132 100644
--- a/magic_buddy.c
+++ b/magic_buddy.c
@@ -251,3 +251,63 @@ void *reallocate(void *old, size_t old_size, size_t new_size,
     }
     return old;
 }
+
+// grow can also be used to move the buddy's data. grow can never fail.
+void grow_buddy(uint8_t *new_base, size_t new_size, struct buddy *state) {
+    // first: make sure all pointers in @state are pointing into @new_base.
+    for (size_t i = 0; i < ADDRESS_BITS; i++) {
+        if (!state->avail[i]) continue;
+        state->avail[i] = (void*)(
+              new_base + ((uint8_t*)state->avail[i] - (uint8_t*)state->base));
+    }
+    state->base = new_base;
+
+    // then, increase the size by 1 repeatedly
+    size_t logsize = size2log(new_size, 0);
+    assert(logsize >= state->root_logsize);
+    if (logsize == state->root_logsize) return;
+    if (isfree((void*)new_base, state->root_logsize, state)) {
+        pop((void*)new_base);
+        push((void*)new_base, state->root_logsize++, state);
+    } else {
+        struct free_block *buddy
+            = buddy_of((void*)new_base, state->root_logsize, state);
+        makefree(buddy, state);
+        push(buddy, state->root_logsize++, state);
+    }
+    return grow_buddy(new_base, new_size, state);
+}
+
+// 0 -> failure
+int shrink_buddy(size_t new_size, struct buddy *state) {
+    // To divide the space in half, we need either:
+    // (1) the root is available, or
+    // (2) the root is split, with the right child being free.
+    // to divide the space by 2^k, we need that property to be true recursively
+    // along the lhs path, until the root itself is free.
+
+    size_t logsize = size2log(new_size, 0);
+
+    // First, just check whether we'll be able to do it:
+    size_t virtual_root_logsize = state->root_logsize;
+    while (virtual_root_logsize > logsize) {
+        if (!isfree(state->base, virtual_root_logsize, state)) {
+            struct free_block *rhs_child
+                = rhs_child_of(state->base, virtual_root_logsize, state);
+            if (!isfree(rhs_child, virtual_root_logsize-1, state))
+                return 0;
+        }
+        virtual_root_logsize--;
+    }
+
+    // It's possible! So go through and actually free the rhs children.
+    while (state->root_logsize > logsize) {
+        if (isfree(state->base, state->root_logsize, state)) break;
+        struct free_block *rhs_child
+            = rhs_child_of(state->base, virtual_root_logsize, state);
+        memset(rhs_child, 0, sizeof(struct free_block));
+        state->root_logsize--;
+    }
+    state->root_logsize = logsize;
+    return 1;
+}
-- 
cgit v1.2.3