1 files changed, 7 insertions, 5 deletions

diff --git a/malloc.c b/malloc.c
index c8d865b..8738373 100644
--- a/malloc.c
+++ b/malloc.c
@@ -5,15 +5,17 @@
 static uint8_t *fseg_base = &edata;
 static uint8_t *malloc_top = &stext;
 
-void *malloc(int n)
+void *malloc_align(int n, int align)
 {
-	malloc_top -= (n + 15) & -16;
+	malloc_top = (uint8_t *) ((uintptr_t)(malloc_top - n) & -align);
 	return malloc_top;
 }
 
-void *malloc_fseg(int n)
+void *malloc_fseg_align(int n, int align)
 {
-	void *p = fseg_base;
-	fseg_base += (n + 15) & -16;
+	void *p;
+	fseg_base = (uint8_t *) (((uintptr_t)fseg_base + align - 1) & -align);
+	p = fseg_base;
+	fseg_base += n;
 	return p;
 }