1 files changed, 148 insertions, 127 deletions

diff --git a/src/color.c b/src/color.c
index f004bf2..7f653b0 100644
--- a/src/color.c
+++ b/src/color.c
@@ -1,9 +1,10 @@
 // Copyright © Tavian Barnes <tavianator@tavianator.com>
 // SPDX-License-Identifier: 0BSD
 
-#include "prelude.h"
 #include "color.h"
+
 #include "alloc.h"
+#include "bfs.h"
 #include "bfstd.h"
 #include "bftw.h"
 #include "diag.h"
@@ -13,6 +14,7 @@
 #include "fsade.h"
 #include "stat.h"
 #include "trie.h"
+
 #include <errno.h>
 #include <fcntl.h>
 #include <stdarg.h>
@@ -161,8 +163,13 @@ static struct esc_seq **get_esc(const struct colors *colors, const char *name) {
 	return leaf ? leaf->value : NULL;
 }
 
+/** Append an escape sequence to a string. */
+static int cat_esc(dchar **dstr, const struct esc_seq *seq) {
+	return dstrxcat(dstr, seq->seq, seq->len);
+}
+
 /** Set a named escape sequence. */
-static int set_esc(struct colors *colors, const char *name, char *value) {
+static int set_esc(struct colors *colors, const char *name, dchar *value) {
 	struct esc_seq **field = get_esc(colors, name);
 	if (!field) {
 		return 0;
@@ -209,57 +216,26 @@ static void ext_tolower(char *ext, size_t len) {
 	}
 }
 
-/**
- * The "smart case" algorithm.
- *
- * @param ext
- *         The current extension being added.
- * @param prev
- *         The previous case-sensitive match, if any, for the same extension.
- * @param iprev
- *         The previous case-insensitive match, if any, for the same extension.
- * @return
- *         Whether this extension should become case-sensitive.
- */
-static bool ext_case_sensitive(struct ext_color *ext, struct ext_color *prev, struct ext_color *iprev) {
-	// This is the first case-insensitive occurrence of this extension, e.g.
-	//
-	//     *.gz=01;31:*.tar.gz=01;33
-	if (!iprev) {
-		bfs_assert(!prev);
-		return false;
-	}
-
-	// If the last version of this extension is already case-sensitive,
-	// this one should be too, e.g.
-	//
-	//     *.tar.gz=01;31:*.TAR.GZ=01;32:*.TAR.GZ=01;33
-	if (iprev->case_sensitive) {
-		return true;
-	}
-
-	// The case matches the last occurrence exactly, e.g.
-	//
-	//     *.tar.gz=01;31:*.tar.gz=01;33
-	if (iprev == prev) {
-		return false;
+/** Insert an extension into a trie. */
+static int insert_ext(struct trie *trie, struct ext_color *ext) {
+	// A later *.x should override any earlier *.x, *.y.x, etc.
+	struct trie_leaf *leaf;
+	while ((leaf = trie_find_postfix(trie, ext->ext))) {
+		trie_remove(trie, leaf);
 	}
 
-	// Different case, but same value, e.g.
-	//
-	//     *.tar.gz=01;31:*.TAR.GZ=01;31
-	if (esc_eq(iprev->esc, ext->esc->seq, ext->esc->len)) {
-		return false;
+	size_t len = ext->len + 1;
+	leaf = trie_insert_mem(trie, ext->ext, len);
+	if (!leaf) {
+		return -1;
 	}
 
-	// Different case, different value, e.g.
-	//
-	//     *.tar.gz=01;31:*.TAR.GZ=01;33
-	return true;
+	leaf->value = ext;
+	return 0;
 }
 
 /** Set the color for an extension. */
-static int set_ext(struct colors *colors, char *key, char *value) {
+static int set_ext(struct colors *colors, dchar *key, dchar *value) {
 	size_t len = dstrlen(key);
 	struct ext_color *ext = varena_alloc(&colors->ext_arena, len + 1);
 	if (!ext) {
@@ -274,45 +250,19 @@ static int set_ext(struct colors *colors, char *key, char *value) {
 		goto fail;
 	}
 
-	key = memcpy(ext->ext, key, len + 1);
+	memcpy(ext->ext, key, len + 1);
 
 	// Reverse the extension (`*.y.x` -> `x.y.*`) so we can use trie_find_prefix()
-	ext_reverse(key, len);
-
-	// Find any pre-existing exact match
-	struct ext_color *prev = NULL;
-	struct trie_leaf *leaf = trie_find_str(&colors->ext_trie, key);
-	if (leaf) {
-		prev = leaf->value;
-		trie_remove(&colors->ext_trie, leaf);
-	}
-
-	// A later *.x should override any earlier *.x, *.y.x, etc.
-	while ((leaf = trie_find_postfix(&colors->ext_trie, key))) {
-		trie_remove(&colors->ext_trie, leaf);
-	}
+	ext_reverse(ext->ext, len);
 
 	// Insert the extension into the case-sensitive trie
-	leaf = trie_insert_str(&colors->ext_trie, key);
-	if (!leaf) {
-		goto fail;
-	}
-	leaf->value = ext;
-
-	// "Smart case": if the same extension is given with two different
-	// capitalizations (e.g. `*.y.x=31:*.Y.Z=32:`), make it case-sensitive
-	ext_tolower(key, len);
-	leaf = trie_insert_str(&colors->iext_trie, key);
-	if (!leaf) {
+	if (insert_ext(&colors->ext_trie, ext) != 0) {
 		goto fail;
 	}
 
-	struct ext_color *iprev = leaf->value;
-	if (ext_case_sensitive(ext, prev, iprev)) {
-		iprev->case_sensitive = true;
-		ext->case_sensitive = true;
+	if (colors->ext_len < len) {
+		colors->ext_len = len;
 	}
-	leaf->value = ext;
 
 	return 0;
 
@@ -324,32 +274,83 @@ fail:
 	return -1;
 }
 
-/** Rebuild the case-insensitive trie after all extensions have been parsed. */
-static int build_iext_trie(struct colors *colors) {
-	trie_clear(&colors->iext_trie);
+/**
+ * The "smart case" algorithm.
+ *
+ * @param ext
+ *         The current extension being added.
+ * @param iext
+ *         The previous case-insensitive match, if any, for the same extension.
+ * @return
+ *         Whether this extension should become case-sensitive.
+ */
+static bool ext_case_sensitive(struct ext_color *ext, struct ext_color *iext) {
+	// This is the first case-insensitive occurrence of this extension, e.g.
+	//
+	//     *.gz=01;31:*.tar.gz=01;33
+	if (!iext) {
+		return false;
+	}
+
+	// If the last version of this extension is already case-sensitive,
+	// this one should be too, e.g.
+	//
+	//     *.tar.gz=01;31:*.TAR.GZ=01;32:*.TAR.GZ=01;33
+	if (iext->case_sensitive) {
+		return true;
+	}
+
+	// Different case, but same value, e.g.
+	//
+	//     *.tar.gz=01;31:*.TAR.GZ=01;31
+	if (esc_eq(iext->esc, ext->esc->seq, ext->esc->len)) {
+		return false;
+	}
+
+	// Different case, different value, e.g.
+	//
+	//     *.tar.gz=01;31:*.TAR.GZ=01;33
+	return true;
+}
 
+/** Build the case-insensitive trie, after all extensions have been parsed. */
+static int build_iext_trie(struct colors *colors) {
+	// Find which extensions should be case-sensitive
 	for_trie (leaf, &colors->ext_trie) {
-		size_t len = leaf->length - 1;
-		if (colors->ext_len < len) {
-			colors->ext_len = len;
+		struct ext_color *ext = leaf->value;
+
+		// "Smart case": if the same extension is given with two different
+		// capitalizations (e.g. `*.y.x=31:*.Y.Z=32:`), make it case-sensitive
+		ext_tolower(ext->ext, ext->len);
+
+		size_t len = ext->len + 1;
+		struct trie_leaf *ileaf = trie_insert_mem(&colors->iext_trie, ext->ext, len);
+		if (!ileaf) {
+			return -1;
 		}
 
+		struct ext_color *iext = ileaf->value;
+		if (ext_case_sensitive(ext, iext)) {
+			ext->case_sensitive = true;
+			iext->case_sensitive = true;
+		}
+
+		ileaf->value = ext;
+	}
+
+	// Rebuild the trie with only the case-insensitive ones
+	trie_clear(&colors->iext_trie);
+
+	for_trie (leaf, &colors->ext_trie) {
 		struct ext_color *ext = leaf->value;
 		if (ext->case_sensitive) {
 			continue;
 		}
 
-		// set_ext() already reversed and lowercased the extension
-		struct trie_leaf *ileaf;
-		while ((ileaf = trie_find_postfix(&colors->iext_trie, ext->ext))) {
-			trie_remove(&colors->iext_trie, ileaf);
-		}
-
-		ileaf = trie_insert_str(&colors->iext_trie, ext->ext);
-		if (!ileaf) {
+		// We already lowercased the extension above
+		if (insert_ext(&colors->iext_trie, ext) != 0) {
 			return -1;
 		}
-		ileaf->value = ext;
 	}
 
 	return 0;
@@ -578,7 +579,7 @@ static int parse_gnu_ls_colors(struct colors *colors, const char *ls_colors) {
 				break;
 			}
 
-			if (dstrncpy(&key, chunk, equals - chunk) != 0) {
+			if (dstrxcpy(&key, chunk, equals - chunk) != 0) {
 				goto fail;
 			}
 			if (unescape(&value, equals + 1, ':', &next) != 0) {
@@ -587,8 +588,8 @@ static int parse_gnu_ls_colors(struct colors *colors, const char *ls_colors) {
 
 			// All-zero values should be treated like NULL, to fall
 			// back on any other relevant coloring for that file
-			char *esc = value;
-			if (strspn(value, "0") == strlen(value)
+			dchar *esc = value;
+			if (strspn(value, "0") == dstrlen(value)
 			    && strcmp(key, "rs") != 0
 			    && strcmp(key, "lc") != 0
 			    && strcmp(key, "rc") != 0
@@ -693,6 +694,20 @@ struct colors *parse_colors(void) {
 		colors->link->len = 0;
 	}
 
+	// Pre-compute the reset escape sequence
+	if (!colors->endcode) {
+		dchar *ec = dstralloc(0);
+		if (!ec
+		    || cat_esc(&ec, colors->leftcode) != 0
+		    || cat_esc(&ec, colors->reset) != 0
+		    || cat_esc(&ec, colors->rightcode) != 0
+		    || set_esc(colors, "ec", ec) != 0) {
+			dstrfree(ec);
+			goto fail;
+		}
+		dstrfree(ec);
+	}
+
 	return colors;
 
 fail:
@@ -727,10 +742,11 @@ CFILE *cfwrap(FILE *file, const struct colors *colors, bool close) {
 	}
 
 	cfile->file = file;
+	cfile->fd = fileno(file);
 	cfile->need_reset = false;
 	cfile->close = close;
 
-	if (isatty(fileno(file))) {
+	if (isatty(cfile->fd)) {
 		cfile->colors = colors;
 	} else {
 		cfile->colors = NULL;
@@ -874,7 +890,7 @@ error:
 
 /** Print an escape sequence chunk. */
 static int print_esc_chunk(CFILE *cfile, const struct esc_seq *esc) {
-	return dstrxcat(&cfile->buffer, esc->seq, esc->len);
+	return cat_esc(&cfile->buffer, esc);
 }
 
 /** Print an ANSI escape sequence. */
@@ -908,12 +924,7 @@ static int print_reset(CFILE *cfile) {
 	}
 	cfile->need_reset = false;
 
-	const struct colors *colors = cfile->colors;
-	if (colors->endcode) {
-		return print_esc_chunk(cfile, colors->endcode);
-	} else {
-		return print_esc(cfile, colors->reset);
-	}
+	return print_esc_chunk(cfile, cfile->colors->endcode);
 }
 
 /** Print a shell-escaped string. */
@@ -953,7 +964,7 @@ static ssize_t first_broken_offset(const char *path, const struct BFTW *ftwbuf,
 	if (path == ftwbuf->path) {
 		if (ftwbuf->depth == 0) {
 			at_fd = AT_FDCWD;
-			at_path = dstrndup(path, max);
+			at_path = dstrxdup(path, max);
 		} else {
 			// The parent must have existed to get here
 			goto out;
@@ -961,14 +972,14 @@ static ssize_t first_broken_offset(const char *path, const struct BFTW *ftwbuf,
 	} else {
 		// We're in print_link_target(), so resolve relative to the link's parent directory
 		at_fd = ftwbuf->at_fd;
-		if (at_fd == AT_FDCWD && path[0] != '/') {
-			at_path = dstrndup(ftwbuf->path, ftwbuf->nameoff);
-			if (at_path && dstrncat(&at_path, path, max) != 0) {
+		if (at_fd == (int)AT_FDCWD && path[0] != '/') {
+			at_path = dstrxdup(ftwbuf->path, ftwbuf->nameoff);
+			if (at_path && dstrxcat(&at_path, path, max) != 0) {
 				ret = -1;
 				goto out_path;
 			}
 		} else {
-			at_path = dstrndup(path, max);
+			at_path = dstrxdup(path, max);
 		}
 	}
 
@@ -1110,7 +1121,7 @@ static int print_link_target(CFILE *cfile, const struct BFTW *ftwbuf) {
 }
 
 /** Format some colored output to the buffer. */
-attr(printf(2, 3))
+_printf(2, 3)
 static int cbuff(CFILE *cfile, const char *format, ...);
 
 /** Dump a parsed expression tree, for debugging. */
@@ -1127,14 +1138,20 @@ static int print_expr(CFILE *cfile, const struct bfs_expr *expr, bool verbose, i
 		return -1;
 	}
 
-	if (bfs_expr_is_parent(expr)) {
-		if (cbuff(cfile, "${red}%pq${rs}", expr->argv[0]) < 0) {
-			return -1;
-		}
-	} else {
-		if (cbuff(cfile, "${blu}%pq${rs}", expr->argv[0]) < 0) {
-			return -1;
-		}
+	int ret;
+	switch (expr->kind) {
+	case BFS_FLAG:
+		ret = cbuff(cfile, "${cyn}%pq${rs}", expr->argv[0]);
+		break;
+	case BFS_OPERATOR:
+		ret = cbuff(cfile, "${red}%pq${rs}", expr->argv[0]);
+		break;
+	default:
+		ret = cbuff(cfile, "${blu}%pq${rs}", expr->argv[0]);
+		break;
+	}
+	if (ret < 0) {
+		return -1;
 	}
 
 	for (size_t i = 1; i < expr->argc; ++i) {
@@ -1156,7 +1173,7 @@ static int print_expr(CFILE *cfile, const struct bfs_expr *expr, bool verbose, i
 	}
 
 	int count = 0;
-	for (struct bfs_expr *child = bfs_expr_children(expr); child; child = child->next) {
+	for_expr (child, expr) {
 		if (dstrcat(&cfile->buffer, " ") != 0) {
 			return -1;
 		}
@@ -1179,10 +1196,9 @@ static int print_expr(CFILE *cfile, const struct bfs_expr *expr, bool verbose, i
 	return 0;
 }
 
-attr(printf(2, 0))
+_printf(2, 0)
 static int cvbuff(CFILE *cfile, const char *format, va_list args) {
 	const struct colors *colors = cfile->colors;
-	int error = errno;
 
 	// Color specifier (e.g. ${blu}) state
 	struct esc_seq **esc;
@@ -1192,7 +1208,7 @@ static int cvbuff(CFILE *cfile, const char *format, va_list args) {
 
 	for (const char *i = format; *i; ++i) {
 		size_t verbatim = strcspn(i, "%$");
-		if (dstrncat(&cfile->buffer, i, verbatim) != 0) {
+		if (dstrxcat(&cfile->buffer, i, verbatim) != 0) {
 			return -1;
 		}
 		i += verbatim;
@@ -1240,12 +1256,6 @@ static int cvbuff(CFILE *cfile, const char *format, va_list args) {
 				}
 				break;
 
-			case 'm':
-				if (dstrcat(&cfile->buffer, xstrerror(error)) != 0) {
-					return -1;
-				}
-				break;
-
 			case 'p':
 				switch (*++i) {
 				case 'q':
@@ -1390,3 +1400,14 @@ int cfprintf(CFILE *cfile, const char *format, ...) {
 	va_end(args);
 	return ret;
 }
+
+int cfreset(CFILE *cfile) {
+	const struct colors *colors = cfile->colors;
+	if (!colors) {
+		return 0;
+	}
+
+	const struct esc_seq *esc = colors->endcode;
+	size_t ret = xwrite(cfile->fd, esc->seq, esc->len);
+	return ret == esc->len ? 0 : -1;
+}