reorganize

2 files changed, 379 insertions, 0 deletions

diff --git a/magic_buddy/buddy.h b/magic_buddy/buddy.h
new file mode 100644
index 0000000..c8778d6
--- /dev/null
+++ b/magic_buddy/buddy.h
@@ -0,0 +1,55 @@
+#pragma once
+#include <stdint.h>
+#include <stddef.h>
+
+#define MAGIC_COOKIE_BYTES 32
+#define ADDRESS_BITS (8 * sizeof(void*))
+
+struct buddy {
+    uint8_t magic[MAGIC_COOKIE_BYTES];
+    struct free_block *(avail[ADDRESS_BITS]);
+    size_t root_logsize;
+    void *base;
+};
+
+// Initialize a buddy with (constant) global state stored in @state.
+// NOTE: after initializing the buddy, you should always pass the exact same
+// pointer in for @state in future calls. If you need to move @state, use
+// move_buddy(...).
+void init_buddy(uint8_t *base, size_t size, uint8_t magic[MAGIC_COOKIE_BYTES],
+                struct buddy *state);
+
+// Allocate a block of size >= @size
+void *allocate(size_t size, struct buddy *state);
+
+// Liberate a block of size @size starting at @base.
+void  liberate(void *base, size_t size, struct buddy *state);
+
+// Print debug information for the allocator.
+void debug_buddy(struct buddy *state);
+
+// Simulates @new = allocate, memcpy(@new, @old), free(@old), with some
+// optimizations for cases where the reallocation can be done in place.
+void *reallocate(void *old, size_t old_size, size_t new_size,
+                 struct buddy *state);
+
+// Attempts to reserve a range [@start,@start+@size).
+// Returns 1 if success, 0 otherwise.
+// Whenever possible, we avoid writing anything into the reserved region.
+int reserve(void *start, size_t size, struct buddy *state);
+
+// Update @state to assume the memory pool has been copied to
+// [@new_base,@new_base+@new_size)
+// Can *ONLY* be used when @new_size >= the existing size.
+void grow_buddy(uint8_t *new_base, size_t new_size, struct buddy *state);
+
+// Update @state to only use the subset of the pool in range
+// [@state->base,@state->base+new_size)
+// Can *ONLY* be used when @new_size <= the existing size.
+// This *CAN* write to anything in the old pool.
+// This *CAN* fail, in which case everything is unmodified and 0 is returned.
+// Upon success, 1 is returned.
+int shrink_buddy(size_t new_size, struct buddy *state);
+
+// Used to move the global state of the buddy.
+void move_buddy(struct buddy *new_state, struct buddy *old_state);
diff --git a/magic_buddy/magic_buddy.c b/magic_buddy/magic_buddy.c
new file mode 100644
index 0000000..80bc069
--- /dev/null
+++ b/magic_buddy/magic_buddy.c
@@ -0,0 +1,324 @@
+#include <stdio.h>
+#include <string.h>
+#include <assert.h>
+#include "buddy.h"
+
+struct free_block {
+    struct free_block *next;
+    struct free_block **pprev;
+    uint8_t magic[MAGIC_COOKIE_BYTES];
+    uint8_t logsize;
+};
+
+// log2 of size rounded up (ceil) or down (!ceil) to the nearest power of 2
+static size_t size2log(size_t size, int ceil) {
+    int bitcnt = 0;
+    for (int i = 0; size; i++) {
+        bitcnt += size & 1;
+        size >>= 1;
+        if (size) continue;
+        if (ceil) return (bitcnt > 1) ? (i + 1) : i;
+        return i;
+    }
+    return ceil ? 1 : 0;
+}
+
+void init_buddy(uint8_t *base, size_t size, uint8_t magic[MAGIC_COOKIE_BYTES],
+                struct buddy *state) {
+    memset(state, 0, sizeof(*state));
+    state->base = base;
+
+    size_t logsize = size2log(size, 0);
+    state->root_logsize = logsize;
+
+    state->avail[logsize] = (struct free_block *)base;
+    state->avail[logsize]->next = 0;
+    state->avail[logsize]->pprev = &(state->avail[logsize]);
+    memcpy(state->magic, magic, sizeof(state->magic));
+}
+
+static struct free_block *buddy_of(struct free_block *block, size_t logsize,
+                                   struct buddy *state) {
+    size_t virt_block = (size_t)block - (size_t)state->base;
+    size_t virt_buddy = virt_block ^ (1 << logsize);
+    return (void*)((uint8_t*)state->base + virt_buddy);
+}
+
+// interpret @block as a block at depth @logsize. is it free?
+static int isfree(struct free_block *block, size_t logsize,
+                  struct buddy *state) {
+    return !memcmp(block->magic, state->magic, sizeof(state->magic))
+        && block->logsize == logsize;
+}
+
+static void makefree(struct free_block *block, struct buddy *state) {
+    memcpy(block->magic, state->magic, sizeof(state->magic));
+}
+
+static struct free_block *pop(struct free_block *block) {
+    *(block->pprev) = block->next;
+    if (block->next) block->next->pprev = block->pprev;
+    return block;
+}
+
+static void push(struct free_block *block, size_t logsize,
+                 struct buddy *state) {
+    block->next = state->avail[logsize];
+    if (block->next) block->next->pprev = &(block->next);
+    state->avail[logsize] = block;
+    block->pprev = &(state->avail[logsize]);
+    block->logsize = logsize;
+}
+
+static void *_allocate(size_t logsize, struct buddy *state) {
+    if (logsize > state->root_logsize) return 0;
+    if (state->avail[logsize]) {
+        struct free_block *block = pop(state->avail[logsize]);
+        memset(block, 0, sizeof(struct free_block));
+        return block;
+    }
+
+    struct free_block *parent = _allocate(logsize + 1, state);
+    struct free_block *buddy = buddy_of(parent, logsize, state);
+    // split @parent in half and place the buddy on the avail list.
+    memcpy(buddy->magic, state->magic, sizeof(state->magic));
+    push(buddy, logsize, state);
+    return parent;
+}
+void *allocate(size_t size, struct buddy *state) {
+    if (size < sizeof(struct free_block)) size = sizeof(struct free_block);
+    return _allocate(size2log(size, 1), state);
+}
+
+static void _liberate(struct free_block *block, size_t logsize,
+                      struct buddy *state) {
+    push(block, logsize, state);
+    if (logsize == state->root_logsize) return;
+
+    struct free_block *buddy = buddy_of(block, logsize, state);
+    if (!isfree(buddy, logsize, state)) return;
+
+    // coalesce up!
+    struct free_block *smaller = (buddy < block) ? buddy : block;
+    struct free_block *bigger = (buddy < block) ? block : buddy;
+    pop(bigger);
+    memset(bigger, 0, sizeof(*bigger));
+
+    pop(smaller);
+
+    _liberate(smaller, logsize + 1, state);
+}
+
+void liberate(void *base, size_t size, struct buddy *state) {
+    struct free_block *block = base;
+    memset(block, 0, sizeof(*block));
+    makefree(block, state);
+
+    _liberate(block, size2log(size, 1), state);
+}
+
+void debug_buddy(struct buddy *state) {
+    for (size_t i = 0; i <= state->root_logsize; i++) {
+        printf("Free blocks at size 2^%lu = %lu:\n", i, 1ul << i);
+        for (struct free_block *block = state->avail[i];
+                block; block = block->next)
+            printf("\t%p\n", block);
+    }
+}
+
+///////// "advanced features"
+static struct free_block *rhs_child_of(struct free_block *block, size_t logsize,
+                                       struct buddy *state) {
+    size_t virt_block = (size_t)block - (size_t)state->base;
+    size_t virt_child = virt_block | (1 << (logsize - 1));
+    return (void*)((uint8_t*)state->base + virt_child);
+}
+
+// NOTE: this method is perhaps more complicated than it needs to be because we
+// 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;
+        if (!parent_logsize--) return 0;
+    }
+    // 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);
+    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) ? parent : rhs_child;
+        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
+            makefree(rhs_child, state);
+            push(rhs_child, parent_logsize - 1, state);
+            parent = lhs_child;
+        }
+        parent_logsize--;
+    }
+    return 0;
+}
+
+static void *_naive_reallocate(void *old, size_t old_size, size_t new_size,
+                               struct buddy *state) {
+    assert(old_size < new_size);
+    void *new = allocate(new_size, state);
+    if (!new) return 0;
+    memcpy(new, old, old_size);
+    liberate(old, old_size, state);
+    return new;
+}
+
+void *reallocate(void *old, size_t old_size, size_t new_size,
+                 struct buddy *state) {
+    if (new_size == 0) return liberate(old, old_size, state), (void*)0;
+
+    if (new_size < sizeof(struct free_block))
+        new_size = sizeof(struct free_block);
+
+    size_t old_logsize = size2log(old_size, 1);
+    size_t new_logsize = size2log(new_size, 1);
+
+    if (new_logsize == old_logsize) return old;
+
+    if (new_logsize < old_logsize) {
+        // repeatedly split, keeping lhs.
+        struct free_block *block = old;
+        while (new_logsize < old_logsize) {
+            old_logsize--;
+            struct free_block *right_half
+                = buddy_of(block, old_logsize, state);
+            makefree(right_half, state);
+            push(right_half, old_logsize, state);
+        }
+        return old;
+    }
+
+    // otherwise, we must iterate up the tree and ensure that at each level:
+    // (1) we are the left-buddy, and
+    // (2) our buddy is free.
+    // up until we reach an ancestor of the proper size.
+
+    // First, just verify this claim.
+    size_t pos_logsize = old_logsize;
+    if (new_logsize > state->root_logsize) return 0;
+    struct free_block *pos = old;
+    while (pos_logsize != new_logsize) {
+        struct free_block *buddy = buddy_of(pos, pos_logsize, state);
+        if (pos > buddy) {
+            // oh no, we're the right buddy!
+            return _naive_reallocate(old, old_size, new_size, state);
+        } else if (!isfree(buddy, pos_logsize, state)) {
+            // oh no, our buddy at this level isn't free!
+            return _naive_reallocate(old, old_size, new_size, state);
+        }
+        pos_logsize++;
+    }
+
+    // Then, revisit the path and coalesce on the way up.
+    pos_logsize = old_logsize;
+    pos = old;
+    while (pos_logsize != new_logsize) {
+        struct free_block *buddy = buddy_of(pos, pos_logsize, state);
+        pop(buddy);
+        memset(buddy, 0, sizeof(struct free_block));
+        pos_logsize++;
+    }
+    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;
+}
+
+void move_buddy(struct buddy *new_state, struct buddy *old_state) {
+    memcpy(new_state, old_state, sizeof(struct buddy));
+    for (size_t i = 0; i < ADDRESS_BITS; i++) {
+        if (!new_state->avail[i]) continue;
+        new_state->avail[i]->pprev = &(new_state->avail[i]);
+    }
+}