From 63cc3d7464446ac9bcca7bdffcbba1380186af7d Mon Sep 17 00:00:00 2001
From: Fabian Freyer <fabian.freyer@physik.tu-berlin.de>
Date: Tue, 20 Feb 2018 15:27:59 +0100
Subject: [PATCH] allocator: add support for page-aligned chunks.

---
 allocator.c            | 48 ++++++++++++++++++++++++++++++++++++++++++
 allocator.h            | 15 +++++++++++++
 tests/test-allocator.c | 44 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 107 insertions(+)

diff --git a/allocator.c b/allocator.c
index e35bc04..73b93bc 100644
--- a/allocator.c
+++ b/allocator.c
@@ -143,6 +143,54 @@ allocate(size_t size)
     return last_endp;
 }
 
+static inline void*
+__align(void* ptr)
+{
+    if ((uintptr_t) ptr % PAGESZ == 0)
+        return ptr;
+    
+    return ptr - ((uintptr_t) ptr % PAGESZ) + PAGESZ;
+}
+
+void*
+allocate_aligned(size_t size)
+{
+    allocation_t *it = NULL, *new_allocation = NULL;
+    void *last_endp = min_alloc_address; /* pointer to the end of the last allocation seen */
+
+    /*
+     * If we do not have any allocations yet, just allocate at the first
+     * possible address.
+     */
+    if (unlikely(TAILQ_EMPTY(&allocations))) {
+        MK_ALLOCATION(new_allocation, __align(min_alloc_address), size);
+        TAILQ_INSERT_HEAD(&allocations, new_allocation, entry);
+        return (void*) min_alloc_address;
+    }
+    
+    /* 
+     * We never remove anything from the list, so we don't need to use the safe
+     * version here.
+     */
+    TAILQ_FOREACH(it, &allocations, entry) {
+        if (last_endp && (it->addr >= last_endp + size)) {
+            /* We found a sufficiently large space. */
+            MK_ALLOCATION(new_allocation, last_endp, size);
+            TAILQ_INSERT_BEFORE(it, new_allocation, entry);
+            return last_endp;
+        }
+        last_endp = __align(it->addr + it->size);
+    }
+
+    /*
+     * We arrived at the end of the list without finding a sufficiently
+     * large free space. Let's allocate one at the end now.
+     */
+    MK_ALLOCATION(new_allocation, last_endp, size);
+    TAILQ_INSERT_TAIL(&allocations, new_allocation, entry);
+    return last_endp;
+}
+
 void*
 allocate_at(void* at, size_t size)
 {
diff --git a/allocator.h b/allocator.h
index 6138f62..9fc1dff 100644
--- a/allocator.h
+++ b/allocator.h
@@ -33,6 +33,8 @@
 #define MiB (size_t) (1024 * 1024)
 #define GiB (size_t) (1024 * 1024 * 1024)
 
+#define PAGESZ ((size_t) 4*kiB) 
+
 /**
  * @brief Initialize the guest memory allocator with the memory mapping
  * information.
@@ -54,6 +56,19 @@ void init_allocator(size_t lowmem, size_t highmem);
  */
 void* allocate(size_t size);
 
+/**
+ * @brief Find a free area for a page-aligned chunk of at least @param size
+ * bytes
+ * 
+ * This is a simple first-fit algorithm.
+ * Iterate over allocations until a fitting space is found, then allocate
+ * a chunk of @param size bytes at the start of this area.
+ * 
+ * @param size size in bytes to allocate
+ * @return void* pointer to the page-aligned allocated chunk
+ */
+void* allocate_aligned(size_t size);
+
 /**
  * @brief Attempt to allocate a chunk of @param size bytes at @param at address.
  * 
diff --git a/tests/test-allocator.c b/tests/test-allocator.c
index 54f83fa..92f3cec 100644
--- a/tests/test-allocator.c
+++ b/tests/test-allocator.c
@@ -251,12 +251,56 @@ ATF_TC_BODY(mmap, tc)
                         p);
 }
 
+ATF_TC(aligned);
+ATF_TC_HEAD(aligned, tc)
+{
+    atf_tc_set_md_var(tc, "descr", "Test alligned allocations");
+}
+ATF_TC_BODY(aligned, tc)
+{
+    void* p = NULL;
+    init_allocator(1 * MiB +0x5000, 0);
+
+    /* Allocate something in the first page */
+    p = allocate_aligned(0x100);
+    ATF_CHECK_EQ_MSG(p, (void*) (1 * MiB),
+                        "Expected allocation at %p, but pointer returned was at"
+                        " %p.",
+                        (void*) (1 * MiB),
+                        p);
+
+    /* Allocate something in the third pages */
+    allocate_at((void*) (1 * MiB) + 0x2020, PAGESZ-0x20);
+
+    p = allocate_aligned(0x100);
+    ATF_CHECK_EQ_MSG(p, (void*) (1 * MiB + 0x1000),
+                        "Expected allocation at %p, but pointer returned was at"
+                        " %p.",
+                        (void*) (1 * MiB +0x1000),
+                        p);
+
+    p = allocate_aligned(0x100);
+    ATF_CHECK_EQ_MSG(p, (void*) (1 * MiB + 0x3000),
+                        "Expected allocation at %p, but pointer returned was at"
+                        " %p.",
+                        (void*) (1 * MiB +0x3000),
+                        p);
+
+    p = allocate_aligned(0x100);
+    ATF_CHECK_EQ_MSG(p, (void*) (1 * MiB + 0x4000),
+                        "Expected allocation at %p, but pointer returned was at"
+                        " %p.",
+                        (void*) (1 * MiB +0x4000),
+                        p);
+}
+
 ATF_TP_ADD_TCS(tp)
 {
     ATF_TP_ADD_TC(tp, sequential);
     ATF_TP_ADD_TC(tp, allocate_at);
     ATF_TP_ADD_TC(tp, firstfit);
     ATF_TP_ADD_TC(tp, mmap);
+    ATF_TP_ADD_TC(tp, aligned);
 
     return atf_no_error();
 }