trie: Make leaves into a linked list

5 files changed, 91 insertions, 45 deletions

diff --git a/src/color.c b/src/color.c
index 6908628..f96813f 100644
--- a/src/color.c
+++ b/src/color.c
@@ -486,17 +486,14 @@ struct colors *parse_colors(void) {
 
 void free_colors(struct colors *colors) {
 	if (colors) {
-		struct trie_leaf *leaf;
-		while ((leaf = trie_first_leaf(&colors->ext_colors))) {
+		TRIE_FOR_EACH(&colors->ext_colors, leaf) {
 			dstrfree(leaf->value);
-			trie_remove(&colors->ext_colors, leaf);
 		}
 		trie_destroy(&colors->ext_colors);
 
-		while ((leaf = trie_first_leaf(&colors->names))) {
+		TRIE_FOR_EACH(&colors->names, leaf) {
 			char **field = leaf->value;
 			dstrfree(*field);
-			trie_remove(&colors->names, leaf);
 		}
 		trie_destroy(&colors->names);
 
diff --git a/src/ctx.c b/src/ctx.c
index 8ba2f38..b2f7d56 100644
--- a/src/ctx.c
+++ b/src/ctx.c
@@ -270,8 +270,7 @@ int bfs_ctx_free(struct bfs_ctx *ctx) {
 		bfs_groups_free(ctx->groups);
 		bfs_users_free(ctx->users);
 
-		struct trie_leaf *leaf;
-		while ((leaf = trie_first_leaf(&ctx->files))) {
+		TRIE_FOR_EACH(&ctx->files, leaf) {
 			struct bfs_ctx_file *ctx_file = leaf->value;
 
 			if (bfs_ctx_fclose(ctx, ctx_file) != 0) {
@@ -282,7 +281,6 @@ int bfs_ctx_free(struct bfs_ctx *ctx) {
 			}
 
 			free(ctx_file);
-			trie_remove(&ctx->files, leaf);
 		}
 		trie_destroy(&ctx->files);
 
diff --git a/src/trie.c b/src/trie.c
index b38ff68..bde1e39 100644
--- a/src/trie.c
+++ b/src/trie.c
@@ -172,6 +172,8 @@ static uintptr_t trie_encode_node(const struct trie_node *node) {
 
 void trie_init(struct trie *trie) {
 	trie->root = 0;
+	trie->head = NULL;
+	trie->tail = NULL;
 }
 
 /** Check if a number is a power of two. */
@@ -242,10 +244,6 @@ static struct trie_leaf *trie_representative(const struct trie *trie, const void
 	return trie_decode_leaf(ptr);
 }
 
-struct trie_leaf *trie_first_leaf(const struct trie *trie) {
-	return trie_representative(trie, NULL, 0);
-}
-
 struct trie_leaf *trie_find_str(const struct trie *trie, const char *key) {
 	return trie_find_mem(trie, key, strlen(key) + 1);
 }
@@ -342,16 +340,47 @@ struct trie_leaf *trie_find_prefix(const struct trie *trie, const char *key) {
 }
 
 /** Create a new leaf, holding a copy of the given key. */
-static struct trie_leaf *new_trie_leaf(const void *key, size_t length) {
+static struct trie_leaf *trie_leaf_alloc(struct trie *trie, const void *key, size_t length) {
 	struct trie_leaf *leaf = malloc(BFS_FLEX_SIZEOF(struct trie_leaf, key, length));
-	if (leaf) {
-		leaf->value = NULL;
-		leaf->length = length;
-		memcpy(leaf->key, key, length);
+	if (!leaf) {
+		return NULL;
+	}
+
+	leaf->prev = trie->tail;
+	leaf->next = NULL;
+
+	if (leaf->prev) {
+		leaf->prev->next = leaf;
+	} else {
+		trie->head = leaf;
 	}
+
+	trie->tail = leaf;
+
+	leaf->value = NULL;
+	leaf->length = length;
+	memcpy(leaf->key, key, length);
+
 	return leaf;
 }
 
+/** Free a leaf. */
+static void trie_leaf_free(struct trie *trie, struct trie_leaf *leaf) {
+	if (leaf->prev) {
+		leaf->prev->next = leaf->next;
+	} else {
+		trie->head = leaf->next;
+	}
+
+	if (leaf->next) {
+		leaf->next->prev = leaf->prev;
+	} else {
+		trie->tail = leaf->prev;
+	}
+
+	free(leaf);
+}
+
 /** Compute the size of a trie node with a certain number of children. */
 static size_t trie_node_size(unsigned int size) {
 	// Empty nodes aren't supported
@@ -410,7 +439,7 @@ static size_t trie_key_mismatch(const void *key1, const void *key2, size_t lengt
  *      | Z
  *      +--->...
  */
-static struct trie_leaf *trie_node_insert(uintptr_t *ptr, const void *key, size_t length, size_t offset) {
+static struct trie_leaf *trie_node_insert(struct trie *trie, uintptr_t *ptr, const void *key, size_t length, size_t offset) {
 	struct trie_node *node = trie_decode_node(*ptr);
 	unsigned int size = trie_popcount(node->bitmap);
 
@@ -423,7 +452,7 @@ static struct trie_leaf *trie_node_insert(uintptr_t *ptr, const void *key, size_
 		*ptr = trie_encode_node(node);
 	}
 
-	struct trie_leaf *leaf = new_trie_leaf(key, length);
+	struct trie_leaf *leaf = trie_leaf_alloc(trie, key, length);
 	if (!leaf) {
 		return NULL;
 	}
@@ -507,7 +536,7 @@ static uintptr_t *trie_jump(uintptr_t *ptr, const char *key, size_t *offset) {
  *      | Y
  *      +--->key
  */
-static struct trie_leaf *trie_split(uintptr_t *ptr, const void *key, size_t length, struct trie_leaf *rep, size_t offset, size_t mismatch) {
+static struct trie_leaf *trie_split(struct trie *trie, uintptr_t *ptr, const void *key, size_t length, struct trie_leaf *rep, size_t offset, size_t mismatch) {
 	unsigned char key_nibble = trie_key_nibble(key, mismatch);
 	unsigned char rep_nibble = trie_key_nibble(rep->key, mismatch);
 	assert(key_nibble != rep_nibble);
@@ -517,7 +546,7 @@ static struct trie_leaf *trie_split(uintptr_t *ptr, const void *key, size_t leng
 		return NULL;
 	}
 
-	struct trie_leaf *leaf = new_trie_leaf(key, length);
+	struct trie_leaf *leaf = trie_leaf_alloc(trie, key, length);
 	if (!leaf) {
 		free(node);
 		return NULL;
@@ -546,7 +575,7 @@ struct trie_leaf *trie_insert_str(struct trie *trie, const char *key) {
 struct trie_leaf *trie_insert_mem(struct trie *trie, const void *key, size_t length) {
 	struct trie_leaf *rep = trie_representative(trie, key, length);
 	if (!rep) {
-		struct trie_leaf *leaf = new_trie_leaf(key, length);
+		struct trie_leaf *leaf = trie_leaf_alloc(trie, key, length);
 		if (leaf) {
 			trie->root = trie_encode_leaf(leaf);
 		}
@@ -576,7 +605,7 @@ struct trie_leaf *trie_insert_mem(struct trie *trie, const void *key, size_t len
 			ptr = &node->children[index];
 		} else {
 			assert(offset == mismatch);
-			return trie_node_insert(ptr, key, length, offset);
+			return trie_node_insert(trie, ptr, key, length, offset);
 		}
 	}
 
@@ -587,11 +616,11 @@ struct trie_leaf *trie_insert_mem(struct trie *trie, const void *key, size_t len
 		}
 	}
 
-	return trie_split(ptr, key, length, rep, offset, mismatch);
+	return trie_split(trie, ptr, key, length, rep, offset, mismatch);
 }
 
 /** Free a chain of singleton nodes. */
-static void trie_free_singletons(uintptr_t ptr) {
+static void trie_free_singletons(struct trie *trie, uintptr_t ptr) {
 	while (!trie_is_leaf(ptr)) {
 		struct trie_node *node = trie_decode_node(ptr);
 
@@ -602,7 +631,7 @@ static void trie_free_singletons(uintptr_t ptr) {
 		free(node);
 	}
 
-	free(trie_decode_leaf(ptr));
+	trie_leaf_free(trie, trie_decode_leaf(ptr));
 }
 
 /**
@@ -667,14 +696,14 @@ void trie_remove(struct trie *trie, struct trie_leaf *leaf) {
 	assert(trie_decode_leaf(*child) == leaf);
 
 	if (!parent) {
-		trie_free_singletons(trie->root);
+		trie_free_singletons(trie, trie->root);
 		trie->root = 0;
 		return;
 	}
 
 	struct trie_node *node = trie_decode_node(*parent);
 	child = node->children + child_index;
-	trie_free_singletons(*child);
+	trie_free_singletons(trie, *child);
 
 	node->bitmap ^= child_bit;
 	unsigned int parent_size = trie_popcount(node->bitmap);
diff --git a/src/trie.h b/src/trie.h
index 2d29ac7..2ede6ea 100644
--- a/src/trie.h
+++ b/src/trie.h
@@ -1,6 +1,6 @@
 /****************************************************************************
  * bfs                                                                      *
- * Copyright (C) 2019 Tavian Barnes <tavianator@tavianator.com>             *
+ * Copyright (C) 2019-2022 Tavian Barnes <tavianator@tavianator.com>        *
  *                                                                          *
  * Permission to use, copy, modify, and/or distribute this software for any *
  * purpose with or without fee is hereby granted.                           *
@@ -17,21 +17,20 @@
 #ifndef BFS_TRIE_H
 #define BFS_TRIE_H
 
+#include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
 
 /**
- * A trie that holds a set of fixed- or variable-length strings.
- */
-struct trie {
-	uintptr_t root;
-};
-
-/**
  * A leaf of a trie.
  */
 struct trie_leaf {
 	/**
+	 * Linked list of leaves, in insertion order.
+	 */
+	struct trie_leaf *prev, *next;
+
+	/**
 	 * An arbitrary value associated with this leaf.
 	 */
 	void *value;
@@ -48,19 +47,19 @@ struct trie_leaf {
 };
 
 /**
- * Initialize an empty trie.
+ * A trie that holds a set of fixed- or variable-length strings.
  */
-void trie_init(struct trie *trie);
+struct trie {
+	/** Pointer to the root node/leaf. */
+	uintptr_t root;
+	/** Linked list of leaves. */
+	struct trie_leaf *head, *tail;
+};
 
 /**
- * Get the first (lexicographically earliest) leaf in the trie.
- *
- * @param trie
- *         The trie to search.
- * @return
- *         The first leaf, or NULL if the trie is empty.
+ * Initialize an empty trie.
  */
-struct trie_leaf *trie_first_leaf(const struct trie *trie);
+void trie_init(struct trie *trie);
 
 /**
  * Find the leaf for a string key.
@@ -153,4 +152,12 @@ void trie_remove(struct trie *trie, struct trie_leaf *leaf);
  */
 void trie_destroy(struct trie *trie);
 
+/**
+ * Iterate over the leaves of a trie.
+ */
+#define TRIE_FOR_EACH(trie, leaf)				\
+	for (struct trie_leaf *leaf = (trie)->head, *_next;	\
+	     leaf && (_next = leaf->next, true);		\
+	     leaf = _next)
+
 #endif // BFS_TRIE_H
diff --git a/tests/trie.c b/tests/trie.c
index 0158fd8..6bc7549 100644
--- a/tests/trie.c
+++ b/tests/trie.c
@@ -82,6 +82,17 @@ int main(void) {
 		assert(leaf->length == strlen(keys[i]) + 1);
 	}
 
+	{
+		size_t i = 0;
+		TRIE_FOR_EACH(&trie, leaf) {
+			assert(leaf == trie_find_str(&trie, keys[i]));
+			assert(!leaf->prev || leaf->prev->next == leaf);
+			assert(!leaf->next || leaf->next->prev == leaf);
+			++i;
+		}
+		assert(i == nkeys);
+	}
+
 	for (size_t i = 0; i < nkeys; ++i) {
 		struct trie_leaf *leaf = trie_find_str(&trie, keys[i]);
 		assert(leaf);
@@ -110,6 +121,10 @@ int main(void) {
 		}
 	}
 
+	TRIE_FOR_EACH(&trie, leaf) {
+		assert(false);
+	}
+
 	// This tests the "jump" node handling on 32-bit platforms
 	size_t longsize = 1 << 20;
 	char *longstr = malloc(longsize);