Add custom allocator support to u8upper()

7 files changed, 49 insertions, 27 deletions

diff --git a/include/_alloc_fn.h b/include/_alloc_fn.h
index 0ff6b58..09a792d 100644
--- a/include/_alloc_fn.h
+++ b/include/_alloc_fn.h
@@ -1,6 +1,7 @@
 #ifndef MLIB__ALLOC_FN_H
 #define MLIB__ALLOC_FN_H
 
-typedef void *(*alloc_fn)(void *ctx, void *ptr, size_t old, size_t new);
+typedef void *(*alloc_fn)(void *ctx, void *ptr, size_t old, size_t new,
+                          size_t align);
 
 #endif /* !MLIB__ALLOC_FN_H */
diff --git a/include/alloc.h b/include/alloc.h
index 57e7f59..966361e 100644
--- a/include/alloc.h
+++ b/include/alloc.h
@@ -20,9 +20,9 @@ typedef struct {
 	size_t _init;
 } arena;
 
+/* Heap allocation functions */
 [[nodiscard, gnu::returns_nonnull]] void *bufalloc(void *, size_t, size_t);
 [[nodiscard]] void *bufalloc_noterm(void *, size_t, size_t);
-[[nodiscard]] void *heapalloc(void *, void *, size_t, size_t);
 
 [[_mlib_pure, _mlib_inline]]
 static inline arena
@@ -44,4 +44,7 @@ void arena_free(arena *);
 #define arena_resz(a, T, p, n)                                                 \
 	((T *)arena_realloc((a), (p), (n), sizeof(T), alignof(T)))
 
+[[nodiscard]] void *alloc_arena(void *, void *, size_t, size_t, size_t);
+[[nodiscard]] void *alloc_heap(void *, void *, size_t, size_t, size_t);
+
 #endif /* !MLIB_ALLOC_H */
diff --git a/include/unicode/string.h b/include/unicode/string.h
index 60eb972..7da4385 100644
--- a/include/unicode/string.h
+++ b/include/unicode/string.h
@@ -3,6 +3,7 @@
 
 #include <stddef.h>
 
+#include "_alloc_fn.h"
 #include "_charN_t.h"
 #include "_u8view.h"
 
@@ -29,14 +30,14 @@ size_t u8gnext(struct u8view *, struct u8view *);
 size_t u8wnext(struct u8view *, struct u8view *);
 size_t u8wnext_human(struct u8view *, struct u8view *);
 
-[[mlib_warn_trunc]] size_t u8casefold(char8_t *restrict, size_t, struct u8view,
-                                      enum caseflags);
-[[mlib_warn_trunc]] size_t u8lower(char8_t *restrict, size_t, struct u8view,
-                                   enum caseflags);
-[[mlib_warn_trunc]] size_t u8title(char8_t *restrict, size_t, struct u8view,
-                                   enum caseflags);
-[[mlib_warn_trunc]] size_t u8upper(char8_t *restrict, size_t, struct u8view,
-                                   enum caseflags);
+[[mlib_warn_trunc]]
+size_t u8casefold(char8_t *restrict, size_t, struct u8view, enum caseflags);
+[[mlib_warn_trunc]]
+size_t u8lower(char8_t *restrict, size_t, struct u8view, enum caseflags);
+[[mlib_warn_trunc]]
+size_t u8title(char8_t *restrict, size_t, struct u8view, enum caseflags);
+[[nodiscard]]
+char8_t *u8upper(size_t *, struct u8view, enum caseflags, alloc_fn, void *);
 
 constexpr double U8LOWER_SCALE = 1.5;
 constexpr double U8LOWER_SCALE_LT = 3;
diff --git a/lib/alloc/alloc_arena.c b/lib/alloc/alloc_arena.c
new file mode 100644
index 0000000..9ddc6e2
--- /dev/null
+++ b/lib/alloc/alloc_arena.c
@@ -0,0 +1,9 @@
+#include <stdlib.h>
+
+#include "alloc.h"
+
+void *
+alloc_arena(void *ctx, void *ptr, size_t old, size_t new, size_t align)
+{
+	return arena_realloc(ctx, ptr, old, new, 1, align);
+}
diff --git a/lib/alloc/heapalloc.c b/lib/alloc/alloc_heap.c
index 0a698f0..72cd7ad 100644
--- a/lib/alloc/heapalloc.c
+++ b/lib/alloc/alloc_heap.c
@@ -3,7 +3,7 @@
 #include "alloc.h"
 
 void *
-heapalloc(void *, void *ptr, size_t, size_t new)
+alloc_heap(void *, void *ptr, size_t, size_t new, size_t)
 {
 	if (new > 0)
 		return realloc(ptr, new);
diff --git a/lib/alloc/arena_realloc.c b/lib/alloc/arena_realloc.c
index f13ce17..b30c1b2 100644
--- a/lib/alloc/arena_realloc.c
+++ b/lib/alloc/arena_realloc.c
@@ -11,6 +11,9 @@ arena_realloc(arena *a, void *ptr, size_t old, size_t new, size_t elemsz,
 {
 	ASSUME(a != nullptr);
 
+	if (ptr == nullptr)
+		return arena_alloc(a, new, elemsz, align);
+
 	if (old == new)
 		return ptr;
 
@@ -55,7 +58,5 @@ arena_realloc(arena *a, void *ptr, size_t old, size_t new, size_t elemsz,
 
 	/* At this point we just make a new allocation and copy the data over */
 	void *dst = arena_alloc(a, new, elemsz, align);
-	if (dst == nullptr || ptr == nullptr)
-		return nullptr;
-	return memcpy(dst, ptr, new * elemsz);
+	return dst == nullptr ? nullptr : memcpy(dst, ptr, new * elemsz);
 }
diff --git a/lib/unicode/string/u8upper.c b/lib/unicode/string/u8upper.c
index 6d4026d..5da49ba 100644
--- a/lib/unicode/string/u8upper.c
+++ b/lib/unicode/string/u8upper.c
@@ -1,10 +1,13 @@
+#include <errno.h>
+#include <stdckdint.h>
+
 #include "mbstring.h"
 #include "unicode/prop.h"
 #include "unicode/string.h"
 
-size_t
-u8upper(char8_t *restrict dst, size_t dstn, struct u8view sv,
-        enum caseflags flags)
+char8_t *
+u8upper(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
+        void *alloc_ctx)
 {
 	struct ucctx ctx = {
 		.az_or_tr = flags & CF_LANG_AZ,
@@ -12,19 +15,22 @@ u8upper(char8_t *restrict dst, size_t dstn, struct u8view sv,
 		.ẞ = flags & CF_ẞ,
 	};
 
+	size_t bufsz;
+	if (ckd_mul(&bufsz, sv.len, (size_t)U8UPPER_SCALE)) {
+		errno = EOVERFLOW;
+		return nullptr;
+	}
+
+	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, alignof(char8_t));
+	if (dst == nullptr)
+		return nullptr;
+
 	rune ch;
 	size_t n = 0;
-
 	while (u8next(&ch, &sv)) {
 		struct rview rv = uprop_get_uc(ch, ctx);
-		for (size_t i = 0; i < rv.len; i++) {
-			if (n >= dstn) {
-				char8_t buf[U8_LEN_MAX];
-				n += rtou8(buf, sizeof(buf), rv.p[i]);
-			} else
-				n += rtou8(dst + n, dstn - n, rv.p[i]);
-		}
-
+		for (size_t i = 0; i < rv.len; i++)
+			n += rtou8(dst + n, bufsz - n, rv.p[i]);
 		if (ctx.lt) {
 			enum uprop_ccc ccc;
 			if (uprop_is_sd(ch))
@@ -34,5 +40,6 @@ u8upper(char8_t *restrict dst, size_t dstn, struct u8view sv,
 		}
 	}
 
-	return n;
+	*dstn = n;
+	return alloc(alloc_ctx, dst, bufsz, n, alignof(char8_t));
 }