From a85565334708840dd44b849a91c1fa09b9038093 Mon Sep 17 00:00:00 2001
From: Tavian Barnes <tavianator@tavianator.com>
Date: Wed, 5 Feb 2025 13:29:05 -0500
Subject: trie: New trie_{get,set}_{str,mem}() functions

---
 src/color.c | 19 +++----------------
 src/mtab.c  | 11 ++++-------
 src/trie.c  | 30 ++++++++++++++++++++++++++++++
 src/trie.h  | 56 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 93 insertions(+), 23 deletions(-)

diff --git a/src/color.c b/src/color.c
index fdaf2eb..0cc950b 100644
--- a/src/color.c
+++ b/src/color.c
@@ -143,13 +143,7 @@ static int init_esc(struct colors *colors, const char *name, const char *value,
 
 	*field = esc;
 
-	struct trie_leaf *leaf = trie_insert_str(&colors->names, name);
-	if (!leaf) {
-		return -1;
-	}
-
-	leaf->value = field;
-	return 0;
+	return trie_set_str(&colors->names, name, field);
 }
 
 /** Check if an escape sequence is equal to a string. */
@@ -159,8 +153,7 @@ static bool esc_eq(const struct esc_seq *esc, const char *str, size_t len) {
 
 /** Get an escape sequence from the table. */
 static struct esc_seq **get_esc(const struct colors *colors, const char *name) {
-	const struct trie_leaf *leaf = trie_find_str(&colors->names, name);
-	return leaf ? leaf->value : NULL;
+	return trie_get_str(&colors->names, name);
 }
 
 /** Append an escape sequence to a string. */
@@ -225,13 +218,7 @@ static int insert_ext(struct trie *trie, struct ext_color *ext) {
 	}
 
 	size_t len = ext->len + 1;
-	leaf = trie_insert_mem(trie, ext->ext, len);
-	if (!leaf) {
-		return -1;
-	}
-
-	leaf->value = ext;
-	return 0;
+	return trie_set_mem(trie, ext->ext, len, ext);
 }
 
 /** Set the color for an extension. */
diff --git a/src/mtab.c b/src/mtab.c
index bf9fc53..4cbb9a6 100644
--- a/src/mtab.c
+++ b/src/mtab.c
@@ -256,10 +256,7 @@ static int bfs_mtab_fill_types(struct bfs_mtab *mtab) {
 			continue;
 		}
 
-		struct trie_leaf *leaf = trie_insert_mem(&mtab->types, &sb.dev, sizeof(sb.dev));
-		if (leaf) {
-			leaf->value = mount->type;
-		} else {
+		if (trie_set_mem(&mtab->types, &sb.dev, sizeof(sb.dev), mount->type) != 0) {
 			goto fail;
 		}
 	}
@@ -282,9 +279,9 @@ const char *bfs_fstype(const struct bfs_mtab *mtab, const struct bfs_stat *statb
 		}
 	}
 
-	const struct trie_leaf *leaf = trie_find_mem(&mtab->types, &statbuf->dev, sizeof(statbuf->dev));
-	if (leaf) {
-		return leaf->value;
+	const char *type = trie_get_mem(&mtab->types, &statbuf->dev, sizeof(statbuf->dev));
+	if (type) {
+		return type;
 	} else {
 		return "unknown";
 	}
diff --git a/src/trie.c b/src/trie.c
index a7498d4..a4829e3 100644
--- a/src/trie.c
+++ b/src/trie.c
@@ -240,6 +240,16 @@ struct trie_leaf *trie_find_mem(const struct trie *trie, const void *key, size_t
 	return trie_find_mem_impl(trie, key, length);
 }
 
+void *trie_get_str(const struct trie *trie, const char *key) {
+	const struct trie_leaf *leaf = trie_find_str(trie, key);
+	return leaf ? leaf->value : NULL;
+}
+
+void *trie_get_mem(const struct trie *trie, const void *key, size_t length) {
+	const struct trie_leaf *leaf = trie_find_mem(trie, key, length);
+	return leaf ? leaf->value : NULL;
+}
+
 _trie_clones
 static struct trie_leaf *trie_find_postfix_impl(const struct trie *trie, const char *key) {
 	size_t length = strlen(key);
@@ -626,6 +636,26 @@ struct trie_leaf *trie_insert_mem(struct trie *trie, const void *key, size_t len
 	return trie_insert_mem_impl(trie, key, length);
 }
 
+int trie_set_str(struct trie *trie, const char *key, const void *value) {
+	struct trie_leaf *leaf = trie_insert_str(trie, key);
+	if (leaf) {
+		leaf->value = (void *)value;
+		return 0;
+	} else {
+		return -1;
+	}
+}
+
+int trie_set_mem(struct trie *trie, const void *key, size_t length, const void *value) {
+	struct trie_leaf *leaf = trie_insert_mem(trie, key, length);
+	if (leaf) {
+		leaf->value = (void *)value;
+		return 0;
+	} else {
+		return -1;
+	}
+}
+
 /** Free a chain of singleton nodes. */
 static void trie_free_singletons(struct trie *trie, uintptr_t ptr) {
 	while (trie_is_node(ptr)) {
diff --git a/src/trie.h b/src/trie.h
index 318a23b..d8cecab 100644
--- a/src/trie.h
+++ b/src/trie.h
@@ -69,6 +69,32 @@ struct trie_leaf *trie_find_str(const struct trie *trie, const char *key);
  */
 struct trie_leaf *trie_find_mem(const struct trie *trie, const void *key, size_t length);
 
+/**
+ * Get the value associated with a string key.
+ *
+ * @trie
+ *         The trie to search.
+ * @key
+ *         The key to look up.
+ * @return
+ *         The found value, or NULL if the key is not present.
+ */
+void *trie_get_str(const struct trie *trie, const char *key);
+
+/**
+ * Get the value associated with a fixed-size key.
+ *
+ * @trie
+ *         The trie to search.
+ * @key
+ *         The key to look up.
+ * @length
+ *         The length of the key in bytes.
+ * @return
+ *         The found value, or NULL if the key is not present.
+ */
+void *trie_get_mem(const struct trie *trie, const void *key, size_t length);
+
 /**
  * Find the shortest leaf that starts with a given key.
  *
@@ -119,6 +145,36 @@ 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);
 
+/**
+ * Set the value for a string key.
+ *
+ * @trie
+ *         The trie to modify.
+ * @key
+ *         The key to insert.
+ * @value
+ *         The value to set.
+ * @return
+ *         0 on success, -1 on error.
+ */
+int trie_set_str(struct trie *trie, const char *key, const void *value);
+
+/**
+ * Set the value for a fixed-size key.
+ *
+ * @trie
+ *         The trie to modify.
+ * @key
+ *         The key to insert.
+ * @length
+ *         The length of the key in bytes.
+ * @value
+ *         The value to set.
+ * @return
+ *         0 on success, -1 on error.
+ */
+int trie_set_mem(struct trie *trie, const void *key, size_t length, const void *value);
+
 /**
  * Remove a leaf from a trie.
  *
-- 
cgit v1.2.3