2 files changed, 105 insertions, 0 deletions

diff --git a/src/alloc.c b/src/alloc.c
index d3923b2..9bc7c43 100644
--- a/src/alloc.c
+++ b/src/alloc.c
@@ -286,6 +286,17 @@ static size_t fenwick_node(const struct slab *slab, size_t i) {
 	return ret;
 }
 
+/** Query the Fenwick tree.  Returns sum(bits[0:N*i]). */
+static size_t fenwick_query(const struct slab *slab, size_t i) {
+	// (i & -i) isolates the least-significant bit in i.
+	// https://en.wikipedia.org/wiki/Fenwick_tree#The_interrogation_tree
+	size_t ret = 0;
+	for (; i > 0; i -= i & -i) {
+		ret += fenwick_node(slab, i);
+	}
+	return ret;
+}
+
 /** Update the Fenwick tree. */
 static void fenwick_update(struct slab *slab, size_t i, ptrdiff_t delta) {
 	size_t length = slab_length(slab);
@@ -321,6 +332,25 @@ static size_t fenwick_search_free(struct slab *slab) {
 	return low;
 }
 
+/** Binary search the Fenwick tree for a specific rank. */
+static size_t fenwick_search_next(struct slab *slab, size_t n) {
+	size_t low = 0;
+	size_t bit = slab_length(slab) + 1;
+	bfs_assert(has_single_bit(bit));
+
+	do {
+		bit >>= 1;
+		size_t mid = low + bit;
+		size_t node = fenwick_node(slab, mid);
+		if (node <= n) {
+			low = mid;
+			n -= node;
+		}
+	} while (bit > 1);
+
+	return low;
+}
+
 /** Get the chunk for a bitmap index. */
 static void *nth_chunk(struct slab *slab, size_t i, size_t j) {
 	bfs_assert(i < slab_length(slab));
@@ -367,6 +397,41 @@ static void slab_free(struct slab *slab, void *ptr) {
 	poison(ptr, slab_size(slab));
 }
 
+void *slab_next(struct slab *slab, void *ptr) {
+	// Find an allocated chunk after the given pointer
+	size_t min = 0;
+	if (ptr) {
+		min = 1 + chunk_index(slab, ptr);
+	}
+
+	size_t i = min / SIZE_WIDTH;
+	size_t j = min % SIZE_WIDTH;
+	size_t length = slab_length(slab);
+	if (i >= length) {
+		return NULL;
+	}
+
+	// Check for a 1 at bit j or higher
+	size_t word = bitmap_word(slab, i);
+	size_t bit = (size_t)1 << j;
+	j = trailing_zeros(word & -bit);
+
+	if (j >= SIZE_WIDTH) {
+		// None in the same word, query the Fenwick tree
+		size_t rank = fenwick_query(slab, i + 1);
+		i = fenwick_search_next(slab, rank);
+		if (i >= length) {
+			return NULL;
+		}
+
+		word = bitmap_word(slab, i);
+		j = trailing_zeros(word);
+		bfs_assert(j < SIZE_WIDTH);
+	}
+
+	return nth_chunk(slab, i, j);
+}
+
 /** Free a whole slab. */
 static void slab_destroy(struct slab *slab) {
 	unpoison(&slab->chunks);
diff --git a/src/alloc.h b/src/alloc.h
index af8b391..4d68111 100644
--- a/src/alloc.h
+++ b/src/alloc.h
@@ -284,6 +284,29 @@ void arena_clear(struct arena *arena);
 void arena_destroy(struct arena *arena);
 
 /**
+ * Loop over every allocated object in an arena.
+ *
+ * @type
+ *         The object type.
+ * @item
+ *         The induction variable name.
+ * @arena
+ *         The arena to loop over.
+ */
+#define for_arena(type, item, arena) \
+	for_arena_ (type, _i_##item, item, (arena))
+
+#define for_arena_(type, i, item, arena) \
+	for (size_t i = 0; i < arena->nslabs; ++i) \
+		for (type *item = NULL; (item = slab_next(arena->slabs[i], item));)
+
+/**
+ * @internal
+ * Get the next allocated object in a slab.
+ */
+void *slab_next(struct slab *slab, void *ptr);
+
+/**
  * An arena allocator for flexibly-sized types.
  */
 struct varena {
@@ -392,4 +415,21 @@ void varena_clear(struct varena *varena);
  */
 void varena_destroy(struct varena *varena);
 
+/**
+ * Loop over every allocated object in a varena.
+ *
+ * @type
+ *         The object type.
+ * @item
+ *         The induction variable name.
+ * @varena
+ *         The varena to loop over.
+ */
+#define for_varena(type, item, varena) \
+	for_varena_ (type, _j_##item, item, (varena))
+
+#define for_varena_(type, j, item, varena) \
+	for (size_t j = 0; j < varena->narenas; ++j) \
+		for_arena (type, item, &varena->arenas[j])
+
 #endif // BFS_ALLOC_H