Accept an element size argument to allocators

8 files changed, 24 insertions, 18 deletions

diff --git a/include/_alloc_fn.h b/include/_alloc_fn.h
index 09a792d..aae4bd0 100644
--- a/include/_alloc_fn.h
+++ b/include/_alloc_fn.h
@@ -2,6 +2,6 @@
 #define MLIB__ALLOC_FN_H
 
 typedef void *(*alloc_fn)(void *ctx, void *ptr, size_t old, size_t new,
-                          size_t align);
+                          size_t elemsz, size_t align);
 
 #endif /* !MLIB__ALLOC_FN_H */
diff --git a/include/alloc.h b/include/alloc.h
index b9e03e2..8004cd0 100644
--- a/include/alloc.h
+++ b/include/alloc.h
@@ -54,7 +54,7 @@ struct heap_ctx {
 	jmp_buf *jmp;
 };
 
-[[nodiscard]] void *alloc_arena(void *, void *, size_t, size_t, size_t);
-[[nodiscard]] void *alloc_heap(void *, void *, size_t, size_t, size_t);
+[[nodiscard]] void *alloc_arena(void *, void *, size_t, size_t, size_t, size_t);
+[[nodiscard]] void *alloc_heap(void *, void *, size_t, size_t, size_t, size_t);
 
 #endif /* !MLIB_ALLOC_H */
diff --git a/lib/alloc/alloc_arena.c b/lib/alloc/alloc_arena.c
index 52d0df9..8d2aa11 100644
--- a/lib/alloc/alloc_arena.c
+++ b/lib/alloc/alloc_arena.c
@@ -5,13 +5,14 @@
 #include "macros.h"
 
 void *
-alloc_arena(void *raw_ctx, void *ptr, size_t old, size_t new, size_t align)
+alloc_arena(void *raw_ctx, void *ptr, size_t old, size_t new, size_t elemsz,
+            size_t align)
 {
 	struct arena_ctx *ctx = raw_ctx;
 	ASSUME(ctx != nullptr);
 	ASSUME(ctx->a != nullptr);
 
-	void *p = arena_realloc(ctx->a, ptr, old, new, 1, align);
+	void *p = arena_realloc(ctx->a, ptr, old, new, elemsz, align);
 	if (new == 0 || p != nullptr)
 		return p;
 
diff --git a/lib/alloc/alloc_heap.c b/lib/alloc/alloc_heap.c
index 46fbdc9..fccfcfc 100644
--- a/lib/alloc/alloc_heap.c
+++ b/lib/alloc/alloc_heap.c
@@ -1,21 +1,26 @@
-#include <stddef.h>
+#include <errno.h>
 #include <setjmp.h>
+#include <stdckdint.h>
+#include <stddef.h>
 #include <stdlib.h>
 
 #include "alloc.h"
 #include "errors.h"
 
 void *
-alloc_heap(void *raw_ctx, void *ptr, size_t, size_t new, size_t)
+alloc_heap(void *raw_ctx, void *ptr, size_t, size_t new, size_t elemsz, size_t)
 {
 	if (new == 0) {
 		free(ptr);
 		return nullptr;
 	}
 
-	void *p = realloc(ptr, new);
-	if (p != nullptr)
-		return p;
+	if (!ckd_mul(&new, new, elemsz)) {
+		void *p = realloc(ptr, new);
+		if (p != nullptr)
+			return p;
+	} else
+		errno = EOVERFLOW;
 
 	struct heap_ctx *ctx = raw_ctx;
 	if (ctx == nullptr || ctx->jmp == nullptr)
diff --git a/lib/unicode/string/u8casefold.c b/lib/unicode/string/u8casefold.c
index b4afe50..c7694a9 100644
--- a/lib/unicode/string/u8casefold.c
+++ b/lib/unicode/string/u8casefold.c
@@ -19,7 +19,7 @@ u8casefold(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 		return nullptr;
 	}
 
-	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, alignof(char8_t));
+	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, 1, alignof(char8_t));
 
 	rune ch;
 	size_t n = 0;
@@ -30,5 +30,5 @@ u8casefold(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 	}
 
 	*dstn = n;
-	return alloc(alloc_ctx, dst, bufsz, n, alignof(char8_t));
+	return alloc(alloc_ctx, dst, bufsz, n, 1, alignof(char8_t));
 }
diff --git a/lib/unicode/string/u8lower.c b/lib/unicode/string/u8lower.c
index 45309b6..cd6a79f 100644
--- a/lib/unicode/string/u8lower.c
+++ b/lib/unicode/string/u8lower.c
@@ -46,7 +46,7 @@ u8lower(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 		return nullptr;
 	}
 
-	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, alignof(char8_t));
+	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, 1, alignof(char8_t));
 
 	while (u8next(&ch, &sv)) {
 		rune next = 0;
@@ -103,5 +103,5 @@ u8lower(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 	}
 
 	*dstn = n;
-	return alloc(alloc_ctx, dst, bufsz, n, alignof(char8_t));
+	return alloc(alloc_ctx, dst, bufsz, n, 1, alignof(char8_t));
 }
diff --git a/lib/unicode/string/u8title.c b/lib/unicode/string/u8title.c
index 5710920..53855b6 100644
--- a/lib/unicode/string/u8title.c
+++ b/lib/unicode/string/u8title.c
@@ -52,7 +52,7 @@ u8title(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 		return nullptr;
 	}
 
-	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, alignof(char8_t));
+	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, 1, alignof(char8_t));
 
 	while (u8next(&ch, &sv)) {
 		if (sv.p > word.p + word.len) {
@@ -134,5 +134,5 @@ u8title(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 	}
 
 	*dstn = n;
-	return alloc(alloc_ctx, dst, bufsz, n, alignof(char8_t));
+	return alloc(alloc_ctx, dst, bufsz, n, 1, alignof(char8_t));
 }
diff --git a/lib/unicode/string/u8upper.c b/lib/unicode/string/u8upper.c
index 91ae679..4c47613 100644
--- a/lib/unicode/string/u8upper.c
+++ b/lib/unicode/string/u8upper.c
@@ -25,7 +25,7 @@ u8upper(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 		return nullptr;
 	}
 
-	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, alignof(char8_t));
+	char8_t *dst = alloc(alloc_ctx, nullptr, 0, bufsz, 1, alignof(char8_t));
 
 	rune ch;
 	size_t n = 0;
@@ -43,5 +43,5 @@ u8upper(size_t *dstn, struct u8view sv, enum caseflags flags, alloc_fn alloc,
 	}
 
 	*dstn = n;
-	return alloc(alloc_ctx, dst, bufsz, n, alignof(char8_t));
+	return alloc(alloc_ctx, dst, bufsz, n, 1, alignof(char8_t));
 }