1 files changed, 117 insertions, 70 deletions

diff --git a/src/dstring.c b/src/dstring.c
index 60a7df9..0f08679 100644
--- a/src/dstring.c
+++ b/src/dstring.c
@@ -2,124 +2,154 @@
 // SPDX-License-Identifier: 0BSD
 
 #include "dstring.h"
+
 #include "alloc.h"
 #include "bit.h"
 #include "diag.h"
+
 #include <stdarg.h>
+#include <stddef.h>
+#include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
 
 /**
- * The memory representation of a dynamic string.  Users get a pointer to data.
+ * The memory representation of a dynamic string.  Users get a pointer to str.
  */
 struct dstring {
-	size_t capacity;
-	size_t length;
-	char data[];
+	/** Capacity of the string, *including* the terminating NUL. */
+	size_t cap;
+	/** Length of the string, *excluding* the terminating NUL. */
+	size_t len;
+	/** The string itself. */
+	alignas(dchar) char str[];
 };
 
-/** Get the string header from the string data pointer. */
-static struct dstring *dstrheader(const char *dstr) {
-	return (struct dstring *)(dstr - offsetof(struct dstring, data));
+#define DSTR_OFFSET offsetof(struct dstring, str)
+
+/** Back up to the header from a pointer to dstring::str. */
+static struct dstring *dstrheader(const dchar *dstr) {
+	return (struct dstring *)(dstr - DSTR_OFFSET);
+}
+
+/**
+ * In some provenance models, the expression `header->str` has its provenance
+ * restricted to just the `str` field itself, making a future dstrheader()
+ * illegal.  This alternative is guaranteed to preserve provenance for the entire
+ * allocation.
+ *
+ * - https://stackoverflow.com/q/25296019
+ * - https://mastodon.social/@void_friend@tech.lgbt/111144859908104311
+ */
+static dchar *dstrdata(struct dstring *header) {
+	return (char *)header + DSTR_OFFSET;
 }
 
-/** Get the correct size for a dstring with the given capacity. */
-static size_t dstrsize(size_t capacity) {
-	return sizeof_flex(struct dstring, data, capacity + 1);
+/** Set the length of a dynamic string. */
+static void dstrsetlen(struct dstring *header, size_t len) {
+	bfs_assert(len < header->cap);
+	header->len = len;
+	header->str[len] = '\0';
 }
 
 /** Allocate a dstring with the given contents. */
-static char *dstralloc_impl(size_t capacity, size_t length, const char *data) {
+static dchar *dstralloc_impl(size_t cap, size_t len, const char *str) {
 	// Avoid reallocations for small strings
-	if (capacity < 7) {
-		capacity = 7;
+	if (cap < DSTR_OFFSET) {
+		cap = DSTR_OFFSET;
 	}
 
-	struct dstring *header = malloc(dstrsize(capacity));
+	struct dstring *header = ALLOC_FLEX(struct dstring, str, cap);
 	if (!header) {
 		return NULL;
 	}
 
-	header->capacity = capacity;
-	header->length = length;
+	header->cap = cap;
+	dstrsetlen(header, len);
 
-	memcpy(header->data, data, length);
-	header->data[length] = '\0';
-	return header->data;
+	dchar *ret = dstrdata(header);
+	memcpy(ret, str, len);
+	return ret;
 }
 
-char *dstralloc(size_t capacity) {
-	return dstralloc_impl(capacity, 0, "");
+dchar *dstralloc(size_t cap) {
+	return dstralloc_impl(cap + 1, 0, "");
 }
 
-char *dstrdup(const char *str) {
+dchar *dstrdup(const char *str) {
 	return dstrxdup(str, strlen(str));
 }
 
-char *dstrndup(const char *str, size_t n) {
+dchar *dstrndup(const char *str, size_t n) {
 	return dstrxdup(str, strnlen(str, n));
 }
 
-char *dstrddup(const char *dstr) {
+dchar *dstrddup(const dchar *dstr) {
 	return dstrxdup(dstr, dstrlen(dstr));
 }
 
-char *dstrxdup(const char *str, size_t len) {
-	return dstralloc_impl(len, len, str);
+dchar *dstrxdup(const char *str, size_t len) {
+	return dstralloc_impl(len + 1, len, str);
 }
 
-size_t dstrlen(const char *dstr) {
-	return dstrheader(dstr)->length;
+size_t dstrlen(const dchar *dstr) {
+	return dstrheader(dstr)->len;
 }
 
-int dstreserve(char **dstr, size_t capacity) {
+int dstreserve(dchar **dstr, size_t cap) {
 	if (!*dstr) {
-		*dstr = dstralloc(capacity);
+		*dstr = dstralloc(cap);
 		return *dstr ? 0 : -1;
 	}
 
 	struct dstring *header = dstrheader(*dstr);
+	size_t old_cap = header->cap;
+	size_t new_cap = cap + 1; // Terminating NUL
+	if (old_cap >= new_cap) {
+		return 0;
+	}
 
-	if (capacity > header->capacity) {
-		capacity = bit_ceil(capacity + 1) - 1;
-
-		header = realloc(header, dstrsize(capacity));
-		if (!header) {
-			return -1;
-		}
-		header->capacity = capacity;
-
-		*dstr = header->data;
+	new_cap = bit_ceil(new_cap);
+	header = REALLOC_FLEX(struct dstring, str, header, old_cap, new_cap);
+	if (!header) {
+		return -1;
 	}
 
+	header->cap = new_cap;
+	*dstr = dstrdata(header);
 	return 0;
 }
 
-int dstresize(char **dstr, size_t length) {
-	if (dstreserve(dstr, length) != 0) {
+int dstresize(dchar **dstr, size_t len) {
+	if (dstreserve(dstr, len) != 0) {
 		return -1;
 	}
 
 	struct dstring *header = dstrheader(*dstr);
-	header->length = length;
-	header->data[length] = '\0';
+	dstrsetlen(header, len);
 	return 0;
 }
 
-int dstrcat(char **dest, const char *src) {
+void dstrshrink(dchar *dstr, size_t len) {
+	struct dstring *header = dstrheader(dstr);
+	bfs_assert(len <= header->len);
+	dstrsetlen(header, len);
+}
+
+int dstrcat(dchar **dest, const char *src) {
 	return dstrxcat(dest, src, strlen(src));
 }
 
-int dstrncat(char **dest, const char *src, size_t n) {
+int dstrncat(dchar **dest, const char *src, size_t n) {
 	return dstrxcat(dest, src, strnlen(src, n));
 }
 
-int dstrdcat(char **dest, const char *src) {
+int dstrdcat(dchar **dest, const dchar *src) {
 	return dstrxcat(dest, src, dstrlen(src));
 }
 
-int dstrxcat(char **dest, const char *src, size_t len) {
+int dstrxcat(dchar **dest, const char *src, size_t len) {
 	size_t oldlen = dstrlen(*dest);
 	size_t newlen = oldlen + len;
 
@@ -131,23 +161,23 @@ int dstrxcat(char **dest, const char *src, size_t len) {
 	return 0;
 }
 
-int dstrapp(char **str, char c) {
+int dstrapp(dchar **str, char c) {
 	return dstrxcat(str, &c, 1);
 }
 
-int dstrcpy(char **dest, const char *src) {
+int dstrcpy(dchar **dest, const char *src) {
 	return dstrxcpy(dest, src, strlen(src));
 }
 
-int dstrncpy(char **dest, const char *src, size_t n) {
+int dstrncpy(dchar **dest, const char *src, size_t n) {
 	return dstrxcpy(dest, src, strnlen(src, n));
 }
 
-int dstrdcpy(char **dest, const char *src) {
+int dstrdcpy(dchar **dest, const dchar *src) {
 	return dstrxcpy(dest, src, dstrlen(src));
 }
 
-int dstrxcpy(char **dest, const char *src, size_t len) {
+int dstrxcpy(dchar **dest, const char *src, size_t len) {
 	if (dstresize(dest, len) != 0) {
 		return -1;
 	}
@@ -156,19 +186,19 @@ int dstrxcpy(char **dest, const char *src, size_t len) {
 	return 0;
 }
 
-char *dstrprintf(const char *format, ...) {
+dchar *dstrprintf(const char *format, ...) {
 	va_list args;
 
 	va_start(args, format);
-	char *str = dstrvprintf(format, args);
+	dchar *str = dstrvprintf(format, args);
 	va_end(args);
 
 	return str;
 }
 
-char *dstrvprintf(const char *format, va_list args) {
+dchar *dstrvprintf(const char *format, va_list args) {
 	// Guess a capacity to try to avoid reallocating
-	char *str = dstralloc(2*strlen(format));
+	dchar *str = dstralloc(2 * strlen(format));
 	if (!str) {
 		return NULL;
 	}
@@ -181,7 +211,7 @@ char *dstrvprintf(const char *format, va_list args) {
 	return str;
 }
 
-int dstrcatf(char **str, const char *format, ...) {
+int dstrcatf(dchar **str, const char *format, ...) {
 	va_list args;
 
 	va_start(args, format);
@@ -191,25 +221,25 @@ int dstrcatf(char **str, const char *format, ...) {
 	return ret;
 }
 
-int dstrvcatf(char **str, const char *format, va_list args) {
+int dstrvcatf(dchar **str, const char *format, va_list args) {
 	// Guess a capacity to try to avoid calling vsnprintf() twice
 	size_t len = dstrlen(*str);
-	dstreserve(str, len + 2*strlen(format));
-	size_t cap = dstrheader(*str)->capacity;
+	dstreserve(str, len + 2 * strlen(format));
+	size_t cap = dstrheader(*str)->cap;
 
 	va_list copy;
 	va_copy(copy, args);
 
 	char *tail = *str + len;
-	int ret = vsnprintf(tail, cap - len + 1, format, args);
+	size_t tail_cap = cap - len;
+	int ret = vsnprintf(tail, tail_cap, format, args);
 	if (ret < 0) {
 		goto fail;
 	}
 
 	size_t tail_len = ret;
-	if (tail_len > cap - len) {
-		cap = len + tail_len;
-		if (dstreserve(str, cap) != 0) {
+	if (tail_len >= tail_cap) {
+		if (dstreserve(str, len + tail_len) != 0) {
 			goto fail;
 		}
 
@@ -223,20 +253,20 @@ int dstrvcatf(char **str, const char *format, va_list args) {
 
 	va_end(copy);
 
-	struct dstring *header = dstrheader(*str);
-	header->length += tail_len;
+	dstrheader(*str)->len += tail_len;
 	return 0;
 
 fail:
+	va_end(copy);
 	*tail = '\0';
 	return -1;
 }
 
-int dstrescat(char **dest, const char *str, enum wesc_flags flags) {
+int dstrescat(dchar **dest, const char *str, enum wesc_flags flags) {
 	return dstrnescat(dest, str, SIZE_MAX, flags);
 }
 
-int dstrnescat(char **dest, const char *str, size_t n, enum wesc_flags flags) {
+int dstrnescat(dchar **dest, const char *str, size_t n, enum wesc_flags flags) {
 	size_t len = *dest ? dstrlen(*dest) : 0;
 
 	// Worst case growth is `ccc...` => $'\xCC\xCC\xCC...'
@@ -254,8 +284,25 @@ int dstrnescat(char **dest, const char *str, size_t n, enum wesc_flags flags) {
 	return dstresize(dest, cur - *dest);
 }
 
-void dstrfree(char *dstr) {
+void dstrfree(dchar *dstr) {
 	if (dstr) {
 		free(dstrheader(dstr));
 	}
 }
+
+dchar *dstrepeat(const char *str, size_t n) {
+	size_t len = strlen(str);
+	dchar *ret = dstralloc(n * len);
+	if (!ret) {
+		return NULL;
+	}
+
+	for (size_t i = 0; i < n; ++i) {
+		if (dstrxcat(&ret, str, len) < 0) {
+			dstrfree(ret);
+			return NULL;
+		}
+	}
+
+	return ret;
+}