alloc: Implement an arena for flexible structs

1 files changed, 118 insertions, 0 deletions

diff --git a/src/alloc.c b/src/alloc.c
index b437975..a6910ce 100644
--- a/src/alloc.c
+++ b/src/alloc.c
@@ -157,5 +157,123 @@ void arena_destroy(struct arena *arena) {
 	for (size_t i = 0; i < arena->nslabs; ++i) {
 		free(arena->slabs[i]);
 	}
+	free(arena->slabs);
 	sanitize_uninit(arena);
 }
+
+void varena_init(struct varena *varena, size_t align, size_t min, size_t offset, size_t size) {
+	varena->align = align;
+	varena->offset = offset;
+	varena->size = size;
+	varena->narenas = 0;
+	varena->arenas = NULL;
+
+	// The smallest size class is at least as many as fit in the smallest
+	// aligned allocation size
+	size_t min_count = (flex_size(align, min, offset, size, 1) - offset + size - 1) / size;
+	varena->shift = bit_width(min_count - 1);
+}
+
+/** Get the size class for the given array length. */
+static size_t varena_size_class(struct varena *varena, size_t count) {
+	// Since powers of two are common array lengths, make them the
+	// (inclusive) upper bound for each size class
+	return bit_width((count - !!count) >> varena->shift);
+}
+
+/** Get the exact size of a flexible struct. */
+static size_t varena_exact_size(const struct varena *varena, size_t count) {
+	return flex_size(varena->align, 0, varena->offset, varena->size, count);
+}
+
+/** Get the arena for the given array length. */
+static struct arena *varena_get(struct varena *varena, size_t count) {
+	size_t i = varena_size_class(varena, count);
+
+	if (i >= varena->narenas) {
+		size_t narenas = i + 1;
+		struct arena *arenas = realloc(varena->arenas, sizeof_array(struct arena, narenas));
+		if (!arenas) {
+			return NULL;
+		}
+
+		for (size_t j = varena->narenas; j < narenas; ++j) {
+			size_t shift = j + varena->shift;
+			size_t size = varena_exact_size(varena, (size_t)1 << shift);
+			arena_init(&arenas[j], varena->align, size);
+		}
+
+		varena->narenas = narenas;
+		varena->arenas = arenas;
+	}
+
+	return &varena->arenas[i];
+}
+
+void *varena_alloc(struct varena *varena, size_t count) {
+	struct arena *arena = varena_get(varena, count);
+	if (!arena) {
+		return NULL;
+	}
+
+	void *ret = arena_alloc(arena);
+	if (!ret) {
+		return NULL;
+	}
+
+	// Tell the sanitizers the exact size of the allocated struct
+	sanitize_free(ret, arena->size);
+	sanitize_alloc(ret, varena_exact_size(varena, count));
+
+	return ret;
+}
+
+void *varena_realloc(struct varena *varena, void *ptr, size_t old_count, size_t new_count) {
+	struct arena *new_arena = varena_get(varena, new_count);
+	struct arena *old_arena = varena_get(varena, old_count);
+	if (!new_arena) {
+		return NULL;
+	}
+
+	size_t new_exact_size = varena_exact_size(varena, new_count);
+	size_t old_exact_size = varena_exact_size(varena, old_count);
+
+	if (new_arena == old_arena) {
+		if (new_count < old_count) {
+			sanitize_free((char *)ptr + new_exact_size, old_exact_size - new_exact_size);
+		} else if (new_count > old_count) {
+			sanitize_alloc((char *)ptr + old_exact_size, new_exact_size - old_exact_size);
+		}
+		return ptr;
+	}
+
+	void *ret = arena_alloc(new_arena);
+	if (!ret) {
+		return NULL;
+	}
+
+	size_t old_size = old_arena->size;
+	sanitize_alloc((char *)ptr + old_exact_size, old_size - old_exact_size);
+
+	size_t new_size = new_arena->size;
+	size_t min_size = new_size < old_size ? new_size : old_size;
+	memcpy(ret, ptr, min_size);
+
+	arena_free(old_arena, ptr);
+	sanitize_free((char *)ret + new_exact_size, new_size - new_exact_size);
+
+	return ret;
+}
+
+void varena_free(struct varena *varena, void *ptr, size_t count) {
+	struct arena *arena = varena_get(varena, count);
+	arena_free(arena, ptr);
+}
+
+void varena_destroy(struct varena *varena) {
+	for (size_t i = 0; i < varena->narenas; ++i) {
+		arena_destroy(&varena->arenas[i]);
+	}
+	free(varena->arenas);
+	sanitize_uninit(varena);
+}