From a1ae7abcff83a8b4ac7409e81e306b275a5250fb Mon Sep 17 00:00:00 2001
From: "Michael D. Lowis" <mike@mdlowis.com>
Date: Sun, 11 Sep 2022 23:48:46 -0400
Subject: [PATCH] added heap

---
 Kernel.c | 49 +++++++++++++++++++++++++++++++++++++------------
 1 file changed, 37 insertions(+), 12 deletions(-)

diff --git a/Kernel.c b/Kernel.c
index bf6beb7..3d56af7 100644
--- a/Kernel.c
+++ b/Kernel.c
@@ -3,6 +3,7 @@
 #include <stdlib.h>
 #include <stdbool.h>
 #include <stdatomic.h>
+#include <string.h>
 #include <unistd.h>
 #include <pthread.h>
 #include <assert.h>
@@ -14,12 +15,17 @@
 
 */
 
-#define COND_WAIT_TIMEOUT_MS 100u
+#define COND_WAIT_TIMEOUT_MS    100u
+#define WORD_SIZE               sizeof(void*)
+#define WORD_ALIGN(x)           (((x) + WORD_SIZE) & ~(WORD_SIZE))
 
 typedef struct Task_T {
     volatile long* stack_top; // top of tasks's stack
     long* stack_base;         // allocated memory for stack
     struct Task_T* next;      // pointer to next task
+    long* heap_beg;
+    long* heap_ptr;
+    long* heap_end;
 } Task_T;
 
 typedef struct {
@@ -203,7 +209,17 @@ void Kernel_Exit(void)
     PickNewTask(false);
 }
 
-static Task_T* CreateTask(void (*task_fn)(void*), void* arg, int stacksize)
+static void* TaskAllocate(Task_T* task, size_t sz)
+{
+    sz = WORD_ALIGN(sz);
+    size_t free_sz = (size_t)task->heap_end - (size_t)task->heap_ptr;
+    assert(sz <= free_sz);
+    void* ptr = task->heap_ptr;
+    task->heap_ptr = (void*)((size_t)task->heap_end + sz);
+    return ptr;
+}
+
+static Task_T* CreateTask(void (*task_fn)(void*), void* arg, long int argsz, long int stacksz, long int heapsz)
 {
     AcquireLock();
     Task_T* task = Dequeue(&DeadQueue);
@@ -214,27 +230,36 @@ static Task_T* CreateTask(void (*task_fn)(void*), void* arg, int stacksize)
         task = calloc(1, sizeof(Task_T));
     }
 
-    /* allocate stack */
-    if (stacksize == 0)
+    /* create the task's heap and copy the argument onto the tasks heap */
+    task->heap_beg = malloc(WORD_ALIGN(heapsz));
+    task->heap_ptr = task->heap_beg;
+    task->heap_end = (void*)((size_t)task->heap_beg + heapsz);
+    void* newarg = TaskAllocate(task, argsz);
+    memcpy(newarg, arg, argsz);
+
+    /* allocate stack (default size of 32768) */
+    if (stacksz == 0)
     {
-        stacksize = 32768;
+        stacksz = 32768;
     }
-    task->stack_base = realloc(task->stack_base, stacksize);
+    task->stack_base = realloc(task->stack_base, stacksz);
 
-    task->stack_top  = &task->stack_base[stacksize/sizeof(long)-1]; // top of stack
+    /* populate the initial context on the stack so SwapTasks works */
+    task->stack_top  = &task->stack_base[stacksz / sizeof(long)-1]; // top of stack
     *(--task->stack_top) = (long)Kernel_Exit; // coroutine cleanup
     *(--task->stack_top) = (long)task_fn;     // user's function to run (rop style!)
-    *(--task->stack_top) = (long)arg;         // user's function argument (rdi)
+    *(--task->stack_top) = (long)newarg;      // user's function argument (rdi)
     for (int saved = 0; saved < 6; saved++)
     {
         *(--task->stack_top) = 0xdeadbeef; // initial values for saved registers
     }
+
     return task;
 }
 
-void Kernel_Spawn(void (*task_fn)(void*), void* arg, int stacksize)
+void Kernel_Spawn(void (*task_fn)(void*), void* arg, long int argsz, long int stacksz, long int heapsz)
 {
-    Task_T* task = CreateTask(task_fn, arg, stacksize);
+    Task_T* task = CreateTask(task_fn, arg, argsz, stacksz, heapsz);
     AcquireLock();
     Enter(task);
     ReleaseLock();
@@ -254,7 +279,7 @@ static void CpuIdle(void* arg)
 static void* CpuMain(void* arg)
 {
     CpuID = (long int)arg;
-    Running[CpuID].idle = CreateTask(CpuIdle, 0, 0);
+    Running[CpuID].idle = CreateTask(CpuIdle, 0, 0, 0, 0);
     Running[CpuID].task = Running[CpuID].idle;
     StartTask(Running[CpuID].idle);
     return NULL; /* unreachable */
@@ -319,7 +344,7 @@ int main(int argc, char** argv)
     {
         int* val = malloc(sizeof(int));
         *val = i;
-        Kernel_Spawn(task1, val, 0);
+        Kernel_Spawn(task1, &i, sizeof(int), 0, sizeof(int));
     }
 
     /* wait for all jobs to be done */
-- 
2.54.0