1 files changed, 26 insertions, 23 deletions

diff --git a/src/malloc/lite_malloc.c b/src/malloc/lite_malloc.c
index 7643fc2c..008549d6 100644
--- a/src/malloc/lite_malloc.c
+++ b/src/malloc/lite_malloc.c
@@ -4,43 +4,46 @@
 #include <errno.h>
 #include "libc.h"
 
-uintptr_t __brk(uintptr_t);
-
 #define ALIGN 16
 
+void *__expand_heap(size_t *);
+
 void *__simple_malloc(size_t n)
 {
-	static uintptr_t cur, brk;
-	uintptr_t base, new;
+	static char *cur, *end;
 	static volatile int lock[2];
-	size_t align=1;
+	size_t align=1, pad;
+	void *p;
 
 	if (!n) n++;
-	if (n > SIZE_MAX/2) goto toobig;
-
 	while (align<n && align<ALIGN)
 		align += align;
-	n = n + align - 1 & -align;
 
 	LOCK(lock);
-	if (!cur) cur = brk = __brk(0)+16;
-	base = cur + align-1 & -align;
-	if (n > SIZE_MAX - PAGE_SIZE - base) goto fail;
-	if (base+n > brk) {
-		new = base+n + PAGE_SIZE-1 & -PAGE_SIZE;
-		if (__brk(new) != new) goto fail;
-		brk = new;
-	}
-	cur = base+n;
-	UNLOCK(lock);
 
-	return (void *)base;
+	pad = -(uintptr_t)cur & align-1;
+
+	if (n <= SIZE_MAX/2 + ALIGN) n += pad;
+
+	if (n > end-cur) {
+		size_t m = n;
+		char *new = __expand_heap(&m);
+		if (!new) {
+			UNLOCK(lock);
+			return 0;
+		}
+		if (new != end) {
+			cur = new;
+			n -= pad;
+			pad = 0;
+		}
+		end = new + m;
+	}
 
-fail:
+	p = cur + pad;
+	cur += n;
 	UNLOCK(lock);
-toobig:
-	errno = ENOMEM;
-	return 0;
+	return p;
 }
 
 weak_alias(__simple_malloc, malloc);