From 6a5c20276a74e39727ac3ebed2079771af03854c Mon Sep 17 00:00:00 2001
From: John Byrd <johnwbyrd@gmail.com>
Date: Sat, 6 Sep 2025 21:15:38 -0700
Subject: [PATCH] New global default stack size setting __set_heap_limit now
 sanity checks against stack; users can  now grab all memory for heap mos-sim
 flushes every output character -- no waiting for output

These changes were necessary to be able to get pi.c to calculate to 50000 digits on the Commodore 64.  Memory allocation had a tendency to walk all over the stack; this puts some sanity checks in place to prevent others from doing the same.  Default behavior is still 4 KB for heap.
---
 mos-platform/common/c/malloc.cc      | 46 +++++++++++++++++++++++++---
 mos-platform/common/c/malloc.s       |  3 ++
 mos-platform/common/include/stdlib.h |  9 ++++--
 utils/sim/mos-sim.c                  |  1 +
 4 files changed, 52 insertions(+), 7 deletions(-)

diff --git a/mos-platform/common/c/malloc.cc b/mos-platform/common/c/malloc.cc
index d666b03e8..72511757b 100644
--- a/mos-platform/common/c/malloc.cc
+++ b/mos-platform/common/c/malloc.cc
@@ -16,6 +16,20 @@
 
 extern char __heap_start;
 extern char __heap_default_limit;
+extern char __stack_reserve_size;
+
+// Get current soft stack pointer value
+__asm__ (
+".text\n"
+".global __get_current_stack_top\n"
+"__get_current_stack_top:\n"
+"  lda __rc0\n"    // Load low byte of soft stack pointer
+"  ldx __rc1\n"    // Load high byte of soft stack pointer  
+"  rts\n"
+);
+
+// Declare the assembly function for C
+extern "C" size_t __get_current_stack_top();
 
 namespace {
 
@@ -179,6 +193,7 @@ FreeChunk *find_fit(size_t size) {
   return nullptr;
 }
 
+"  rts\n"
 // Allocate at chunk of size bytes from a free chunk. The pointer returned
 // points to the contents (past the chunk header).
 void *allocate_free_chunk(FreeChunk *free_chunk, size_t size) {
@@ -221,24 +236,44 @@ extern "C" {
 
 size_t __heap_limit() { return heap_limit; }
 
-void __set_heap_limit(size_t new_limit) {
+size_t __get_heap_max_safe_size() {
+  size_t stack_location = __get_current_stack_top();
+  size_t reserve_size = (size_t)&__stack_reserve_size;
+  size_t heap_start = (size_t)&__heap_start;
+  
+  if (stack_location <= heap_start + reserve_size) {
+    // Stack is too close to heap start; no safe heap possible
+    return 0;
+  }
+  
+  return stack_location - heap_start - reserve_size;
+}
+
+size_t __set_heap_limit(size_t new_limit) {
   TRACE("__set_heap_limit(%u)\n", new_limit);
 
   // Chunk sizes must be a multiple of two.
   if (new_limit & 1)
     --new_limit;
 
+  // Cap the request to avoid stack collision
+  size_t max_safe = __get_heap_max_safe_size();
+  if (new_limit > max_safe) {
+    TRACE("Capping requested limit %u to max safe %u\n", new_limit, max_safe);
+    new_limit = max_safe;
+  }
+
   if (!initialized) {
     heap_limit = (new_limit < MIN_CHUNK_SIZE) ? MIN_CHUNK_SIZE : new_limit;
     TRACE("Heap not yet initialized. Set limit to %u.\n", heap_limit);
-    return;
+    return heap_limit;
   }
 
   // TODO: We can make this actually shrink the heap too...
   if (new_limit <= heap_limit) {
-    TRACE("New limit %u smaller than current %u; returning.", new_limit,
+    TRACE("New limit %u smaller than current %u; returning.\n", new_limit,
           heap_limit);
-    return;
+    return heap_limit;
   }
 
   size_t grow = new_limit - heap_limit;
@@ -255,7 +290,7 @@ void __set_heap_limit(size_t new_limit) {
     TRACE("Last chunk not free.\n");
     if (grow < MIN_CHUNK_SIZE) {
       TRACE("Not enough new size for a chunk; returning.\n");
-      return;
+      return heap_limit;
     }
     TRACE("Inserting new chunk.\n");
     FreeChunk::insert(heap_end(), grow);
@@ -263,6 +298,7 @@ void __set_heap_limit(size_t new_limit) {
   }
 
   heap_limit = new_limit;
+  return heap_limit;
 }
 
 size_t __heap_bytes_used() { return heap_limit - free_size; }
diff --git a/mos-platform/common/c/malloc.s b/mos-platform/common/c/malloc.s
index e76af396c..963ac2de7 100644
--- a/mos-platform/common/c/malloc.s
+++ b/mos-platform/common/c/malloc.s
@@ -1,2 +1,5 @@
 .weak __heap_default_limit
 __heap_default_limit = 4096
+
+.weak __stack_reserve_size
+__stack_reserve_size = 1024
diff --git a/mos-platform/common/include/stdlib.h b/mos-platform/common/include/stdlib.h
index 426f29f66..4fe8df207 100644
--- a/mos-platform/common/include/stdlib.h
+++ b/mos-platform/common/include/stdlib.h
@@ -150,8 +150,13 @@ size_t __heap_limit();
 
 /* Set the maximum size of the heap.  Note the limitations above. */
 /* Setting the heap limit implicitly allocates the heap.  Don't call this
-   function if you aren't going to use the heap. */
-void __set_heap_limit(size_t limit);
+ * function if you aren't going to use the heap.  Returns the value actually
+ * set -- this value is capped to avoid trashing the stack. 
+ */
+size_t __set_heap_limit(size_t limit);
+
+/* Return the maximum safe heap size to avoid stack collision. */
+size_t __get_heap_max_safe_size(void);
 
 /* Return heap bytes in use, including overhead for heap data structures in
    the existing allocations. */
diff --git a/utils/sim/mos-sim.c b/utils/sim/mos-sim.c
index 6d517ad77..fb4233e9d 100644
--- a/utils/sim/mos-sim.c
+++ b/utils/sim/mos-sim.c
@@ -97,6 +97,7 @@ void write6502(uint16_t address, uint8_t value) {
     exit(value);
   case 0xFFF9:
     putchar(value);
+    fflush(stdout);
     break;
   }
 }