alloc: Don't require the old size in varena_realloc()

Instead, just look up which arena contains the pointer.

8 files changed, 69 insertions, 43 deletions

diff --git a/src/alloc.c b/src/alloc.c
index ef9f6ab..f52b701 100644
--- a/src/alloc.c
+++ b/src/alloc.c
@@ -172,9 +172,8 @@ void arena_init(struct arena *arena, size_t align, size_t size) {
 	arena->size = size;
 }
 
-/** Allocate a new slab. */
-_cold
-static int slab_alloc(struct arena *arena) {
+/** Get the size of the ith slab. */
+static size_t slab_size(const struct arena *arena, size_t i) {
 	// Make the initial allocation size ~4K
 	size_t size = 4096;
 	if (size < arena->size) {
@@ -183,7 +182,14 @@ static int slab_alloc(struct arena *arena) {
 	// Trim off the excess
 	size -= size % arena->size;
 	// Double the size for every slab
-	size <<= arena->nslabs;
+	size <<= i;
+	return size;
+}
+
+/** Allocate a new slab. */
+_cold
+static int slab_alloc(struct arena *arena) {
+	size_t size = slab_size(arena, arena->nslabs);
 
 	// Allocate the slab
 	void *slab = zalloc(arena->align, size);
@@ -224,7 +230,24 @@ void *arena_alloc(struct arena *arena) {
 	return chunk;
 }
 
+/** Check if a pointer comes from this arena. */
+static bool arena_contains(const struct arena *arena, void *ptr) {
+	uintptr_t addr = (uintptr_t)ptr;
+
+	for (size_t i = 0; i < arena->nslabs; ++i) {
+		uintptr_t start = (uintptr_t)arena->slabs[i];
+		uintptr_t end = start + slab_size(arena, i);
+		if (addr >= start && addr < end) {
+			return true;
+		}
+	}
+
+	return false;
+}
+
 void arena_free(struct arena *arena, void *ptr) {
+	bfs_assert(arena_contains(arena, ptr));
+
 	union chunk *chunk = ptr;
 	chunk_set_next(arena, chunk, arena->chunks);
 	arena->chunks = chunk;
@@ -292,6 +315,18 @@ static struct arena *varena_get(struct varena *varena, size_t count) {
 	return &varena->arenas[i];
 }
 
+/** Get the arena containing a given pointer. */
+static struct arena *varena_find(struct varena *varena, void *ptr) {
+	for (size_t i = 0; i < varena->narenas; ++i) {
+		struct arena *arena = &varena->arenas[i];
+		if (arena_contains(arena, ptr)) {
+			return arena;
+		}
+	}
+
+	bfs_abort("No arena contains %p", ptr);
+}
+
 void *varena_alloc(struct varena *varena, size_t count) {
 	struct arena *arena = varena_get(varena, count);
 	if (!arena) {
@@ -310,26 +345,23 @@ void *varena_alloc(struct varena *varena, size_t 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);
+void *varena_realloc(struct varena *varena, void *ptr, size_t count) {
+	struct arena *new_arena = varena_get(varena, count);
+	struct arena *old_arena = varena_find(varena, ptr);
 	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);
+	size_t new_size = new_arena->size;
+	size_t new_exact_size = varena_exact_size(varena, count);
 
+	void *ret;
 	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;
+		ret = ptr;
+		goto done;
 	}
 
-	void *ret = arena_alloc(new_arena);
+	ret = arena_alloc(new_arena);
 	if (!ret) {
 		return NULL;
 	}
@@ -337,12 +369,11 @@ void *varena_realloc(struct varena *varena, void *ptr, size_t old_count, size_t
 	size_t old_size = old_arena->size;
 	sanitize_alloc(ptr, old_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);
-
+done:
 	sanitize_free(ret, new_size);
 	sanitize_alloc(ret, new_exact_size);
 	return ret;
@@ -356,15 +387,15 @@ void *varena_grow(struct varena *varena, void *ptr, size_t *count) {
 	size_t new_shift = varena_size_class(varena, old_count + 1) + varena->shift;
 	size_t new_count = (size_t)1 << new_shift;
 
-	ptr = varena_realloc(varena, ptr, old_count, new_count);
+	ptr = varena_realloc(varena, ptr, new_count);
 	if (ptr) {
 		*count = new_count;
 	}
 	return ptr;
 }
 
-void varena_free(struct varena *varena, void *ptr, size_t count) {
-	struct arena *arena = varena_get(varena, count);
+void varena_free(struct varena *varena, void *ptr) {
+	struct arena *arena = varena_find(varena, ptr);
 	arena_free(arena, ptr);
 }
 
diff --git a/src/alloc.h b/src/alloc.h
index 1fafbab..c4badc4 100644
--- a/src/alloc.h
+++ b/src/alloc.h
@@ -338,10 +338,8 @@ void varena_init(struct varena *varena, size_t align, size_t offset, size_t size
  *         The that allocated the object.
  * @ptr
  *         The object to free.
- * @count
- *         The length of the flexible array.
  */
-void varena_free(struct varena *varena, void *ptr, size_t count);
+void varena_free(struct varena *varena, void *ptr);
 
 /**
  * Arena-allocate a flexible struct.
@@ -363,15 +361,13 @@ void *varena_alloc(struct varena *varena, size_t count);
  *         The varena to allocate from.
  * @ptr
  *         The object to resize.
- * @old_count
- *         The old array length.
- * @new_count
+ * @count
  *         The new array length.
  * @return
  *         The resized struct, or NULL on failure.
  */
 _nodiscard
-void *varena_realloc(struct varena *varena, void *ptr, size_t old_count, size_t new_count);
+void *varena_realloc(struct varena *varena, void *ptr, size_t count);
 
 /**
  * Grow a flexible struct by an arbitrary amount.
@@ -380,7 +376,7 @@ void *varena_realloc(struct varena *varena, void *ptr, size_t old_count, size_t
  *         The varena to allocate from.
  * @ptr
  *         The object to resize.
- * @count
+ * @count[in,out]
  *         Pointer to the flexible array length.
  * @return
  *         The resized struct, or NULL on failure.
diff --git a/src/bftw.c b/src/bftw.c
index 61193d5..c99fb8d 100644
--- a/src/bftw.c
+++ b/src/bftw.c
@@ -796,7 +796,7 @@ static void bftw_file_free(struct bftw_cache *cache, struct bftw_file *file) {
 
 	bftw_stat_recycle(cache, file);
 
-	varena_free(&cache->files, file, file->namelen + 1);
+	varena_free(&cache->files, file);
 }
 
 /**
diff --git a/src/color.c b/src/color.c
index e7f0973..63f1039 100644
--- a/src/color.c
+++ b/src/color.c
@@ -128,7 +128,7 @@ static struct esc_seq *new_esc(struct colors *colors, const char *seq, size_t le
 
 /** Free an escape sequence. */
 static void free_esc(struct colors *colors, struct esc_seq *seq) {
-	varena_free(&colors->esc_arena, seq, seq->len);
+	varena_free(&colors->esc_arena, seq);
 }
 
 /** Initialize a color in the table. */
@@ -280,7 +280,7 @@ fail:
 	if (ext->esc) {
 		free_esc(colors, ext->esc);
 	}
-	varena_free(&colors->ext_arena, ext, len + 1);
+	varena_free(&colors->ext_arena, ext);
 	return -1;
 }
 
diff --git a/src/mtab.c b/src/mtab.c
index bf9fc53..7297421 100644
--- a/src/mtab.c
+++ b/src/mtab.c
@@ -101,7 +101,7 @@ static int bfs_mtab_add(struct bfs_mtab *mtab, const char *path, const char *typ
 shrink:
 	--mtab->nmounts;
 free:
-	varena_free(&mtab->varena, mount, size);
+	varena_free(&mtab->varena, mount);
 	return -1;
 }
 
diff --git a/src/pwcache.c b/src/pwcache.c
index fa19dad..3f63633 100644
--- a/src/pwcache.c
+++ b/src/pwcache.c
@@ -52,7 +52,7 @@ static void *bfs_getent(bfs_getent_fn *fn, const void *key, struct trie_leaf *le
 		}
 	}
 
-	varena_free(varena, ptr, bufsize);
+	varena_free(varena, ptr);
 	return NULL;
 }
 
diff --git a/src/trie.c b/src/trie.c
index a7498d4..dd96414 100644
--- a/src/trie.c
+++ b/src/trie.c
@@ -352,7 +352,7 @@ static struct trie_leaf *trie_leaf_alloc(struct trie *trie, const void *key, siz
 /** Free a leaf. */
 static void trie_leaf_free(struct trie *trie, struct trie_leaf *leaf) {
 	LIST_REMOVE(trie, leaf);
-	varena_free(&trie->leaves, leaf, leaf->length);
+	varena_free(&trie->leaves, leaf);
 }
 
 /** Create a new node. */
@@ -362,16 +362,15 @@ static struct trie_node *trie_node_alloc(struct trie *trie, size_t size) {
 }
 
 /** Reallocate a trie node. */
-static struct trie_node *trie_node_realloc(struct trie *trie, struct trie_node *node, size_t old_size, size_t new_size) {
-	bfs_assert(has_single_bit(old_size));
-	bfs_assert(has_single_bit(new_size));
-	return varena_realloc(&trie->nodes, node, old_size, new_size);
+static struct trie_node *trie_node_realloc(struct trie *trie, struct trie_node *node, size_t size) {
+	bfs_assert(has_single_bit(size));
+	return varena_realloc(&trie->nodes, node, size);
 }
 
 /** Free a node. */
 static void trie_node_free(struct trie *trie, struct trie_node *node, size_t size) {
 	bfs_assert(size == (size_t)count_ones(node->bitmap));
-	varena_free(&trie->nodes, node, size);
+	varena_free(&trie->nodes, node);
 }
 
 #if ENDIAN_NATIVE == ENDIAN_LITTLE
@@ -450,7 +449,7 @@ static struct trie_leaf *trie_node_insert(struct trie *trie, uintptr_t *ptr, str
 
 	// Double the capacity every power of two
 	if (has_single_bit(size)) {
-		node = trie_node_realloc(trie, node, size, 2 * size);
+		node = trie_node_realloc(trie, node, 2 * size);
 		if (!node) {
 			trie_leaf_free(trie, leaf);
 			return NULL;
@@ -724,7 +723,7 @@ static void trie_remove_impl(struct trie *trie, struct trie_leaf *leaf) {
 	--parent_size;
 
 	if (has_single_bit(parent_size)) {
-		node = trie_node_realloc(trie, node, 2 * parent_size, parent_size);
+		node = trie_node_realloc(trie, node, parent_size);
 		if (node) {
 			*parent = trie_encode_node(node);
 		}
diff --git a/tests/alloc.c b/tests/alloc.c
index 4aae515..5eb4e9e 100644
--- a/tests/alloc.c
+++ b/tests/alloc.c
@@ -17,7 +17,7 @@ struct flexible {
 
 /** Check varena_realloc() poisoning for a size combination. */
 static struct flexible *check_varena_realloc(struct varena *varena, struct flexible *flexy, size_t old_count, size_t new_count) {
-	flexy = varena_realloc(varena, flexy, old_count, new_count);
+	flexy = varena_realloc(varena, flexy, new_count);
 	bfs_everify(flexy);
 
 	for (size_t i = 0; i < new_count; ++i) {