alloc: Implement an arena for flexible structs

2 files changed, 210 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);
+}
diff --git a/src/alloc.h b/src/alloc.h
index 65edb92..c2ea09b 100644
--- a/src/alloc.h
+++ b/src/alloc.h
@@ -190,4 +190,96 @@ void arena_free(struct arena *arena, void *ptr);
  */
 void arena_destroy(struct arena *arena);
 
+/**
+ * An arena allocator for flexibly-sized types.
+ */
+struct varena {
+	/** The alignment of the struct. */
+	size_t align;
+	/** The offset of the flexible array. */
+	size_t offset;
+	/** The size of the flexible array elements. */
+	size_t size;
+	/** Shift amount for the smallest size class. */
+	size_t shift;
+	/** The number of arenas of different sizes. */
+	size_t narenas;
+	/** The array of differently-sized arenas. */
+	struct arena *arenas;
+};
+
+/**
+ * Initialize a varena for a struct with the given layout.
+ *
+ * @param varena
+ *         The varena to initialize.
+ * @param align
+ *         alignof(type)
+ * @param min
+ *         sizeof(type)
+ * @param offset
+ *         offsetof(type, flexible_array)
+ * @param size
+ *         sizeof(flexible_array[i])
+ */
+void varena_init(struct varena *varena, size_t align, size_t min, size_t offset, size_t size);
+
+/**
+ * Initialize a varena for the given type and flexible array.
+ *
+ * @param varena
+ *         The varena to initialize.
+ * @param type
+ *         A struct type containing a flexible array.
+ * @param member
+ *         The name of the flexible array member.
+ */
+#define VARENA_INIT(varena, type, member) \
+	varena_init(varena, alignof(type), sizeof(type), offsetof(type, member), sizeof_member(type, member[0]))
+
+/**
+ * Arena-allocate a flexible struct.
+ *
+ * @param varena
+ *         The varena to allocate from.
+ * @param count
+ *         The length of the flexible array.
+ * @return
+ *         The allocated struct, or NULL on failure.
+ */
+void *varena_alloc(struct varena *varena, size_t count);
+
+/**
+ * Resize a flexible struct.
+ *
+ * @param varena
+ *         The varena to allocate from.
+ * @param ptr
+ *         The object to resize.
+ * @param old_count
+ *         The old array lenth.
+ * @param new_count
+ *         The new array length.
+ * @return
+ *         The resized struct, or NULL on failure.
+ */
+void *varena_realloc(struct varena *varena, void *ptr, size_t old_count, size_t new_count);
+
+/**
+ * Free an arena-allocated flexible struct.
+ *
+ * @param varena
+ *         The that allocated the object.
+ * @param ptr
+ *         The object to free.
+ * @param count
+ *         The length of the flexible array.
+ */
+void varena_free(struct varena *varena, void *ptr, size_t count);
+
+/**
+ * Destroy a varena, freeing all allocations.
+ */
+void varena_destroy(struct varena *varena);
+
 #endif // BFS_ALLOC_H