Grow the stack as needed.

This is untuned for performance, but seems to be doing the right thing.

src/vm/wren_compiler.c

@@ -324,6 +324,15 @@ struct sCompiler
   int scopeDepth;
   
   // The current number of slots (locals and temporaries) in use.
+  //
+  // We use this and maxSlots to track the maximum number of additional slots
+  // a function may need while executing. When the function is called, the
+  // fiber will check to ensure its stack has enough room to cover that worst
+  // case and grow the stack if needed.
+  //
+  // This value here doesn't include parameters to the function. Since those
+  // are already pushed onto the stack by the caller and tracked there, we
+  // don't need to double count them here.
   int numSlots;
   
   // The maximum number of slots (locals and temporaries) in use at one time in
@@ -1611,11 +1620,6 @@ static bool finishBlock(Compiler* compiler)
 // In that case, this adds the code to ensure it returns `this`.
 static void finishBody(Compiler* compiler, bool isInitializer)
 {
-  // Now that the parameter list has been compiled, we know how many slots they
-  // use.
-  compiler->numSlots = compiler->numLocals;
-  compiler->maxSlots = compiler->numLocals;
-  
   bool isExpressionBody = finishBlock(compiler);
 
   if (isInitializer)
@@ -3075,8 +3079,6 @@ static void createConstructor(Compiler* compiler, Signature* signature,
 {
   Compiler methodCompiler;
   initCompiler(&methodCompiler, compiler->parser, compiler, false);
-  methodCompiler.numSlots = signature->arity + 1;
-  methodCompiler.maxSlots = methodCompiler.numSlots;
   
   // Allocate the instance.
   emitOp(&methodCompiler, compiler->enclosingClass->isForeign

src/vm/wren_debug.c

@@ -18,7 +18,7 @@ void wrenDebugPrintStackTrace(ObjFiber* fiber)
   for (int i = fiber->numFrames - 1; i >= 0; i--)
   {
     CallFrame* frame = &fiber->frames[i];
-    ObjFn* fn = wrenGetFrameFunction(frame);
+    ObjFn* fn = wrenUpwrapClosure(frame->fn);
 
     // The built-in core module has no name. We explicitly omit it from stack
     // traces since we don't want to highlight to a user the implementation

src/vm/wren_utils.c

@@ -175,3 +175,17 @@ int wrenUtf8DecodeNumBytes(uint8_t byte)
   if ((byte & 0xe0) == 0xc0) return 2;
   return 1;
 }
+
+// From: http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2Float
+int wrenPowerOf2Ceil(int n)
+{
+  n--;
+  n |= n >> 1;
+  n |= n >> 2;
+  n |= n >> 4;
+  n |= n >> 8;
+  n |= n >> 16;
+  n++;
+  
+  return n;
+}

src/vm/wren_utils.h

@@ -49,12 +49,7 @@ typedef struct {
     { \
       if (buffer->capacity < buffer->count + count) \
       { \
-        int capacity = buffer->capacity; \
+        int capacity = wrenPowerOf2Ceil(buffer->count + count); \
-        while (capacity < buffer->count + count) \
-        { \
-          capacity = capacity == 0 ? 8 : capacity * 2; \
-        } \
-        \
         buffer->data = (type*)wrenReallocate(vm, buffer->data, \
             buffer->capacity * sizeof(type), capacity * sizeof(type)); \
         buffer->capacity = capacity; \
@@ -120,4 +115,7 @@ int wrenUtf8Decode(const uint8_t* bytes, uint32_t length);
 // returns 0.
 int wrenUtf8DecodeNumBytes(uint8_t byte);
 
+// Returns the smallest power of two that is equal to or greater than [n].
+int wrenPowerOf2Ceil(int n);
+
 #endif

src/vm/wren_value.c

@@ -147,14 +147,19 @@ ObjClosure* wrenNewClosure(WrenVM* vm, ObjFn* fn)
 
 ObjFiber* wrenNewFiber(WrenVM* vm, Obj* fn)
 {
-  // Allocate the call frames before the fiber in case it triggers a GC.
+  // Allocate the arrays before the fiber in case it triggers a GC.
   CallFrame* frames = ALLOCATE_ARRAY(vm, CallFrame, INITIAL_CALL_FRAMES);
   
+  int stackCapacity = wrenPowerOf2Ceil(wrenUpwrapClosure(fn)->maxSlots);
+  Value* stack = ALLOCATE_ARRAY(vm, Value, stackCapacity);
+  
   ObjFiber* fiber = ALLOCATE(vm, ObjFiber);
   initObj(vm, &fiber->obj, OBJ_FIBER, vm->fiberClass);
   fiber->id = vm->nextFiberId++;
   fiber->frames = frames;
   fiber->frameCapacity = INITIAL_CALL_FRAMES;
+  fiber->stack = stack;
+  fiber->stackCapacity = stackCapacity;
   wrenResetFiber(vm, fiber, fn);
 
   return fiber;
@@ -960,6 +965,8 @@ static void blackenFiber(WrenVM* vm, ObjFiber* fiber)
 
   // Keep track of how much memory is still in use.
   vm->bytesAllocated += sizeof(ObjFiber);
+  vm->bytesAllocated += fiber->frameCapacity * sizeof(CallFrame);
+  vm->bytesAllocated += fiber->stackCapacity * sizeof(Value);
 }
 
 static void blackenFn(WrenVM* vm, ObjFn* fn)
@@ -1117,6 +1124,14 @@ void wrenFreeObj(WrenVM* vm, Obj* obj)
       wrenMethodBufferClear(vm, &((ObjClass*)obj)->methods);
       break;
 
+    case OBJ_FIBER:
+    {
+      ObjFiber* fiber = (ObjFiber*)obj;
+      DEALLOCATE(vm, fiber->frames);
+      DEALLOCATE(vm, fiber->stack);
+      break;
+    }
+      
     case OBJ_FN:
     {
       ObjFn* fn = (ObjFn*)obj;
@@ -1147,7 +1162,6 @@ void wrenFreeObj(WrenVM* vm, Obj* obj)
 
     case OBJ_STRING:
     case OBJ_CLOSURE:
-    case OBJ_FIBER:
     case OBJ_INSTANCE:
     case OBJ_RANGE:
     case OBJ_UPVALUE:

src/vm/wren_value.h

@@ -41,9 +41,6 @@
 // The representation is controlled by the `WREN_NAN_TAGGING` define. If that's
 // defined, Nan tagging is used.
 
-// TODO: Make this externally controllable.
-#define STACK_SIZE 1024
-
 // These macros cast a Value to one of the specific value types. These do *not*
 // perform any validation, so must only be used after the Value has been
 // ensured to be the right type.
@@ -210,9 +207,18 @@ typedef struct
 typedef struct sObjFiber
 {
   Obj obj;
-  Value stack[STACK_SIZE];
+  
+  // The stack of value slots. This is used for holding local variables and
+  // temporaries while the fiber is executing. It heap-allocated and grown as
+  // needed.
+  Value* stack;
+  
+  // A pointer to one past the top-most value on the stack.
   Value* stackTop;
 
+  // The number of allocated slots in the stack array.
+  int stackCapacity;
+  
   // The stack of call frames. This is a dynamic array that grows as needed but
   // never shrinks.
   CallFrame* frames;
@@ -629,14 +635,16 @@ ObjFiber* wrenNewFiber(WrenVM* vm, Obj* fn);
 // Resets [fiber] back to an initial state where it is ready to invoke [fn].
 void wrenResetFiber(WrenVM* vm, ObjFiber* fiber, Obj* fn);
 
-static inline ObjFn* wrenGetFrameFunction(CallFrame* frame)
+// Returns the base ObjFn backing an ObjClosure, or [function] if it already is
+// a raw ObjFn.
+static inline ObjFn* wrenUpwrapClosure(Obj* function)
 {
-  switch (frame->fn->type)
+  switch (function->type)
   {
-    case OBJ_FN:      return (ObjFn*)frame->fn;
+    case OBJ_FN:      return (ObjFn*)function;
-    case OBJ_CLOSURE: return ((ObjClosure*)frame->fn)->fn;
+    case OBJ_CLOSURE: return ((ObjClosure*)function)->fn;
     default:
-      UNREACHABLE();
+      ASSERT(false, "Function should be an ObjFn or ObjClosure.");
       return NULL;
   }
 }
@@ -652,7 +660,7 @@ static inline void wrenAppendCallFrame(WrenVM* vm, ObjFiber* fiber,
   CallFrame* frame = &fiber->frames[fiber->numFrames++];
   frame->stackStart = stackStart;
   frame->fn = function;
-  frame->ip = wrenGetFrameFunction(frame)->bytecode;
+  frame->ip = wrenUpwrapClosure(frame->fn)->bytecode;
 }
 
 ObjForeign* wrenNewForeign(WrenVM* vm, ObjClass* classObj, size_t size);

src/vm/wren_vm.c

@@ -432,6 +432,7 @@ static bool checkArity(WrenVM* vm, Value value, int numArgs)
 static inline void callFunction(
     WrenVM* vm, ObjFiber* fiber, Obj* function, int numArgs)
 {
+  // Grow the call frame array if needed.
   if (fiber->numFrames + 1 > fiber->frameCapacity)
   {
     int max = fiber->frameCapacity * 2;
@@ -441,8 +442,43 @@ static inline void callFunction(
     fiber->frameCapacity = max;
   }
   
-  // TODO: Check for stack overflow. We handle the call frame array growing,
+  // Grow the stack if needed.
-  // but not the stack itself.
+  int stackSize = (int)(fiber->stackTop - fiber->stack);
+  int needed = stackSize + wrenUpwrapClosure(function)->maxSlots;
+  
+  if (fiber->stackCapacity < needed)
+  {
+    int capacity = wrenPowerOf2Ceil(needed);
+    
+    Value* oldStack = fiber->stack;
+    fiber->stack = (Value*)wrenReallocate(vm, fiber->stack,
+        sizeof(Value) * fiber->stackCapacity,
+        sizeof(Value) * capacity);
+    fiber->stackCapacity = capacity;
+    
+    // If the reallocation moves the stack, then we need to shift every pointer
+    // into the stack to point to its new location.
+    if (fiber->stack != oldStack)
+    {
+      // Top of the stack.
+      long offset = fiber->stack - oldStack;
+      fiber->stackTop += offset;
+      
+      // Stack pointer for each call frame.
+      for (int i = 0; i < fiber->numFrames; i++)
+      {
+        fiber->frames[i].stackStart += offset;
+      }
+      
+      // Open upvalues.
+      for (ObjUpvalue* upvalue = fiber->openUpvalues;
+           upvalue != NULL;
+           upvalue = upvalue->next)
+      {
+        upvalue->value += offset;
+      }
+    }
+  }
 
   wrenAppendCallFrame(vm, fiber, function, fiber->stackTop - numArgs);
 }
@@ -744,7 +780,7 @@ static WrenInterpretResult runInterpreter(WrenVM* vm, register ObjFiber* fiber)
       frame = &fiber->frames[fiber->numFrames - 1];    \
       stackStart = frame->stackStart;                  \
       ip = frame->ip;                                  \
-      fn = wrenGetFrameFunction(frame);
+      fn = wrenUpwrapClosure(frame->fn);
 
   // Terminates the current fiber with error string [error]. If another calling
   // fiber is willing to catch the error, transfers control to it, otherwise