Add strict variants of stack and heap allocators

4 files changed, 74 insertions, 0 deletions

diff --git a/include/allocator.h b/include/allocator.h
index 4cec7da..20dbede 100644
--- a/include/allocator.h
+++ b/include/allocator.h
@@ -9,4 +9,6 @@ struct allocator {
 };
 
 extern struct allocator stdlib_stack_allocator;
+extern struct allocator stdlib_strict_stack_allocator;
 extern struct allocator stdlib_heap_allocator;
+extern struct allocator stdlib_strict_heap_allocator;
diff --git a/liballoc.api b/liballoc.api
index 8756c21..d23401f 100644
--- a/liballoc.api
+++ b/liballoc.api
@@ -1,2 +1,4 @@
 stdlib_stack_allocator
+stdlib_strict_stack_allocator
 stdlib_heap_allocator
+stdlib_strict_heap_allocator
diff --git a/liballoc.lds b/liballoc.lds
index 789ea95..25400b8 100644
--- a/liballoc.lds
+++ b/liballoc.lds
@@ -1,6 +1,8 @@
 LIBALLOC_1.1 {
 	global:
 		stdlib_stack_allocator;
+		stdlib_strict_stack_allocator;
 		stdlib_heap_allocator;
+		stdlib_strict_heap_allocator;
 	local: *;
 };
diff --git a/src/allocator.c b/src/allocator.c
index d0f358f..eb07e50 100644
--- a/src/allocator.c
+++ b/src/allocator.c
@@ -7,10 +7,16 @@
 static int stdlib_stack_alloc(void **p, size_t size);
 static int stdlib_stack_realloc(void **p, size_t size);
 static int stdlib_stack_free(void *p);
+static int stdlib_strict_stack_alloc(void **p, size_t size);
+static int stdlib_strict_stack_realloc(void **p, size_t size);
+static int stdlib_strict_stack_free(void *p);
 
 static int stdlib_heap_alloc(void **p, size_t size);
 static int stdlib_heap_realloc(void **p, size_t size);
 static int stdlib_heap_free(void *p);
+static int stdlib_strict_heap_alloc(void **p, size_t size);
+static int stdlib_strict_heap_realloc(void **p, size_t size);
+static int stdlib_strict_heap_free(void *p);
 
 struct allocator stdlib_stack_allocator = {
 	.alloc = stdlib_stack_alloc,
@@ -18,12 +24,24 @@ struct allocator stdlib_stack_allocator = {
 	.free = stdlib_stack_free
 };
 
+struct allocator stdlib_strict_stack_allocator = {
+	.alloc = stdlib_strict_stack_alloc,
+	.realloc = stdlib_strict_stack_realloc,
+	.free = stdlib_strict_stack_free
+};
+
 struct allocator stdlib_heap_allocator = {
 	.alloc = stdlib_heap_alloc,
 	.realloc = stdlib_heap_realloc,
 	.free = stdlib_heap_free
 };
 
+struct allocator stdlib_strict_heap_allocator = {
+	.alloc = stdlib_strict_heap_alloc,
+	.realloc = stdlib_strict_heap_realloc,
+	.free = stdlib_strict_heap_free
+};
+
 int
 stdlib_stack_alloc(void **p, size_t size)
 {
@@ -47,6 +65,27 @@ stdlib_stack_free(void *p)
 }
 
 int
+stdlib_strict_stack_alloc(void **p, size_t size)
+{
+	*p = alloca(size);
+	if (!*p) abort();
+
+	return 0;
+}
+
+int
+stdlib_strict_stack_realloc(void **p, size_t size)
+{
+	abort();
+}
+
+int
+stdlib_strict_stack_free(void *p)
+{
+	return 0;
+}
+
+int
 stdlib_heap_alloc(void **p, size_t size)
 {
 	*p = malloc(size);
@@ -73,3 +112,32 @@ stdlib_heap_free(void *p)
 
 	return 0;
 }
+
+int
+stdlib_strict_heap_alloc(void **p, size_t size)
+{
+	*p = malloc(size);
+	if (!*p) abort();
+
+	return 0;
+}
+
+int
+stdlib_strict_heap_realloc(void **p, size_t size)
+{
+	void *np;
+
+	np = realloc(*p, size);
+	if (!np) abort();
+	*p = np;
+
+	return 0;
+}
+
+int
+stdlib_strict_heap_free(void *p)
+{
+	free(p);
+
+	return 0;
+}