Implement arena allocation resizing

4 files changed, 71 insertions, 2 deletions

diff --git a/include/alloc.h b/include/alloc.h
index 3074b27..57e7f59 100644
--- a/include/alloc.h
+++ b/include/alloc.h
@@ -11,7 +11,7 @@
 
 struct _region {
 	size_t len, cap;
-	void *data;
+	void *data, *last;
 	struct _region *next;
 };
 
@@ -31,12 +31,17 @@ mkarena(size_t n)
 	return (arena){._init = n ? n : MLIB_ARENA_BLKSIZE};
 }
 
+/* Arena allocation functions */
 [[nodiscard, gnu::malloc, gnu::alloc_size(2, 3), gnu::alloc_align(4)]]
 void *arena_alloc(arena *, size_t, size_t, size_t);
-
+[[nodiscard]]
+void *arena_realloc(arena *, void *, size_t, size_t, size_t, size_t);
 void arena_zero(arena *);
 void arena_free(arena *);
 
+/* Arena allocation macro wrappers */
 #define arena_new(a, T, n) ((T *)arena_alloc((a), (n), sizeof(T), alignof(T)))
+#define arena_resz(a, T, p, n)                                                 \
+	((T *)arena_realloc((a), (p), (n), sizeof(T), alignof(T)))
 
 #endif /* !MLIB_ALLOC_H */
diff --git a/lib/alloc/arena_alloc.c b/lib/alloc/arena_alloc.c
index 814b0bb..65123d4 100644
--- a/lib/alloc/arena_alloc.c
+++ b/lib/alloc/arena_alloc.c
@@ -36,6 +36,7 @@ mkregion(size_t cap)
 		errno = save;
 		return nullptr;
 	}
+	r->last = r->data;
 	return r;
 }
 
@@ -63,6 +64,7 @@ arena_alloc(arena *a, size_t sz, size_t n, size_t align)
 		if (nlen <= r->cap) {
 			void *ret = (char *)r->data + off;
 			r->len = nlen;
+			r->last = ret;
 			return ret;
 		}
 	}
diff --git a/lib/alloc/arena_realloc.c b/lib/alloc/arena_realloc.c
new file mode 100644
index 0000000..f13ce17
--- /dev/null
+++ b/lib/alloc/arena_realloc.c
@@ -0,0 +1,61 @@
+#include <errno.h>
+#include <stdckdint.h>
+#include <string.h>
+
+#include "alloc.h"
+#include "macros.h"
+
+void *
+arena_realloc(arena *a, void *ptr, size_t old, size_t new, size_t elemsz,
+              size_t align)
+{
+	ASSUME(a != nullptr);
+
+	if (old == new)
+		return ptr;
+
+	struct _region *cur = a->_head;
+	while (cur != nullptr) {
+		if (ptr >= cur->data && (char *)ptr < (char *)cur->data + cur->cap)
+			break;
+		cur = cur->next;
+	}
+	if (cur == nullptr)
+		cur = a->_head;
+
+	/* cur now points to the region containing ‘ptr’ */
+
+	/* If we are shrinking the buffer, then we don’t need to move our allocation
+	   and can just return it directly.  As a minor optimization if the
+	   allocation we are shrinking is the last allocation in the current region,
+	   we can decrease the region length to make space for more future
+	   allocations. */
+	if (old > new) {
+		if (ptr == cur->last)
+			cur->len -= (old - new) * elemsz;
+		return ptr;
+	}
+
+	ASSUME(old < new);
+
+	/* If we need to grow the given allocation, but it was the last allocation
+	   made in a region, then we first see if we can just eat more trailing free
+	   space in the region to avoid a memcpy(). */
+	if (ptr == cur->last) {
+		size_t need, free = cur->cap - cur->len;
+		if (ckd_mul(&need, new - old, elemsz)) {
+			errno = EOVERFLOW;
+			return nullptr;
+		}
+		if (need <= free) {
+			cur->len += need;
+			return ptr;
+		}
+	}
+
+	/* At this point we just make a new allocation and copy the data over */
+	void *dst = arena_alloc(a, new, elemsz, align);
+	if (dst == nullptr || ptr == nullptr)
+		return nullptr;
+	return memcpy(dst, ptr, new * elemsz);
+}
diff --git a/lib/alloc/arena_zero.c b/lib/alloc/arena_zero.c
index b0e3fbe..296007e 100644
--- a/lib/alloc/arena_zero.c
+++ b/lib/alloc/arena_zero.c
@@ -9,6 +9,7 @@ arena_zero(arena *a)
 	struct _region *cur = a->_head;
 	while (cur != nullptr) {
 		cur->len = 0;
+		cur->last = cur->data;
 		cur = cur->next;
 	}
 }