Manage locals and calls using the stack directly

MoritzGoeckel · MoritzGoeckel · commit fd95c4269284 · 2026-01-04T02:12:33.000+01:00
diff --git a/TODO.md b/TODO.md
@@ -1,7 +1,6 @@
 # Roadmap
 
 - [ ] Garbage collection for heap objects (structs)
-- [ ] Remove locals and add registers: src/executer/ByteCode.h:117
 - [ ] Dynamic Strings
 - [ ] Arrays
 
@@ -21,4 +20,5 @@
 - [ ] FFI for Windows
 
 - [x] Throw when finding a type annotation that does not name a type (struct or primitive)
-- [x] Throw when encountering a function that does not return on all paths
+- [x] Throw when encountering a function that does not return on all paths
+- [x] Remove locals and add registers: src/executer/ByteCode.h:117
diff --git a/mfiles/method_variable_num_locals_0.m b/mfiles/method_variable_num_locals_0.m
@@ -0,0 +1,10 @@
+# expect_result=1
+let f(x) = {
+    if(x > 20) {
+        let z = 1;
+        ret z;
+    }
+    let y = 2;
+    ret y;
+};
+ret f(30);
diff --git a/mfiles/method_variable_num_locals_1.m b/mfiles/method_variable_num_locals_1.m
@@ -0,0 +1,10 @@
+# expect_result=2
+let f(x) = {
+    if(x > 20) {
+        let z = 1;
+        ret z;
+    }
+    let y = 2;
+    ret y;
+};
+ret f(10);
diff --git a/src/emitter/ByteCodeEmitter.cpp b/src/emitter/ByteCodeEmitter.cpp
@@ -29,6 +29,7 @@ void ByteCodeEmitter::run() {
         for(const auto& param : fn.second->getHead()->getParameters()) {
             localNames.push_back(param->getName());
         }
+        num_params = localNames.size();  // Remember parameter count
         function_idxs[fn.first] = code().size();
         process(fn.second->getBody(), false);
     }
@@ -206,7 +207,9 @@ void ByteCodeEmitter::process(const std::shared_ptr<AST::Node>& node, bool hasCo
             if (ret->getExpr()) {
                 process(ret->getExpr(), true);
             }
-            code().push_back(executor::Instruction(executor::Op::RET));
+            // RET arg1=num_params, arg2=num_locals
+            size_t num_locals = localNames.size() - num_params;
+            code().push_back(executor::Instruction(executor::Op::RET, num_params, num_locals));
             break;
         }
         case AST::NodeType::Assign: {
diff --git a/src/emitter/ByteCodeEmitter.h b/src/emitter/ByteCodeEmitter.h
@@ -27,6 +27,7 @@ class ByteCodeEmitter {
     };
     std::vector<Backpatch> backpatches;
     std::vector<std::string> localNames; // Id is idx
+    size_t num_params;  // Number of parameters for current function
 
     public:
     ByteCodeEmitter(const std::map<std::string, std::shared_ptr<AST::Function>> &functions);
diff --git a/src/executer/ByteCode.cpp b/src/executer/ByteCode.cpp
@@ -20,7 +20,7 @@ std::string instructionsToString(const std::vector<Instruction>& instructions, b
         { Op::LOCALS, { "LOCALS", { "ID" } } },
         { Op::LOCALL, { "LOCALL", { "ID" } } },
         { Op::CALL, { "CALL", { "NUM_ARGS" } } },
-        { Op::RET, { "RET", { } } },
+        { Op::RET, { "RET", { "NUM_PARAMS", "NUM_LOCALS" } } },
         { Op::PUSH, { "PUSH", { "VALUE" } } },
         { Op::POP, { "POP", {} } },
         { Op::ADD, { "ADD", {} } },
@@ -108,47 +108,91 @@ ProgramState ByteCodeVM::run(size_t maxInstructions) {
                 break;
             }
             case Op::CALL: {
-                auto return_address = idx;
+                word_t num_params = inst.arg1;
+                word_t jmp_dest = stack.pop();
 
-                auto dest = stack.pop();
-                idx = dest; // Jump to function address
-
-                // Bring parameters to locals
-                std::vector<word_t> locals;
-                for (int i = 0; i < inst.arg1; ++i) {
-                    locals.push_back(stack.pop());
+                // Save parameters temporarily (they're on stack in order: param0, param1, ...)
+                std::vector<word_t> params;
+                for (word_t i = 0; i < num_params; ++i) {
+                    params.push_back(stack.pop());
                 }
 
-                callstack.push_back({locals, return_address});
+                stack.push(idx); // Push return address
+                stack.push(function_stack_base); // Push prev function_stack_base
+                function_stack_base = stack.size(); // Update function_stack_base to current top
+
+                // Push parameters in reverse order (so param0 is at function_stack_base+0)
+                for (auto it = params.rbegin(); it != params.rend(); ++it) {
+                    stack.push(*it);
+                }
+                
+                idx = jmp_dest; // Jump to function
                 break;
             }
             case Op::RET: {
-                if(callstack.empty()){
+                word_t num_params = inst.arg1;
+                word_t num_locals = inst.arg2;
+
+                //                                           function_stack_base
+                //                                                   |
+                // Stack layout: [ret_addr][prev_function_stack_base]^[params...][locals...][return_value?]
+
+                word_t expected_stack_size = function_stack_base + num_params + num_locals;
+                bool has_return_value = false;
+                if (stack.size() > expected_stack_size) {
+                    return_value = stack.pop();
+                    has_return_value = true;
+                } else {
+                    return_value = 0;
+                }
+
+                // Pop locals
+                while (stack.size() > function_stack_base + num_params) {
+                    stack.pop();
+                }
+
+                // Pop parameters
+                for (word_t i = 0; i < num_params; ++i) {
+                    stack.pop();
+                }
+
+                // Check if this is the main function return
+                if (function_stack_base == 0) {
                     return ProgramState::Finished;
                 }
-                idx = callstack.back().return_address;
-                callstack.pop_back();
+
+                function_stack_base = stack.pop(); // Restore function_stack_base
+                idx = stack.pop(); // Jump back to return address
+
+                if (has_return_value) {
+                    stack.push(return_value);
+                }
+
                 break;
             }
             case Op::LOCALS: {
-                // LOCALS ID
-                auto& locals = callstack.back().locals;
+                // LOCALS n: Store stack top into local variable/parameter n
+                word_t value = stack.pop();
 
-                auto idx = inst.arg1 + 1;
-                if(locals.size() <= idx) {
-                    locals.resize(idx); // TODO: Expensive!
+                word_t localIndex = function_stack_base + inst.arg1;
+
+                // Expand stack if necessary
+                while (stack.size() <= localIndex) {
+                    stack.push(0);
                 }
 
-                auto value = stack.pop();
-                locals[inst.arg1] = value; // Store value in local variable
+                stack.set(localIndex, value);
                 break;
             }
             case Op::LOCALL: {
-                // LOCALL ID
-                ASSURE(!callstack.empty(), "ByteCodeVM: Empty callstack");
-                auto& locals = callstack.back().locals;
-                ASSURE(inst.arg1 < locals.size(), "ByteCodeVM: Local variable index out of bounds");
-                stack.push(locals[inst.arg1]); // Push local variable onto stack
+                // LOCALL n: Load local variable/parameter n onto stack
+                word_t localIndex = function_stack_base + inst.arg1;
+
+                ASSURE(localIndex < stack.size(),
+                       "ByteCodeVM: Local variable index out of bounds");
+
+                word_t value = stack.get(localIndex);
+                stack.push(value);
                 break;
             }
             case Op::PRINTS: {
@@ -325,11 +369,12 @@ ProgramState ByteCodeVM::run(size_t maxInstructions) {
     return ProgramState::Paused;
 }
 
-ByteCodeVM::ByteCodeVM(const Program& program) : 
-    idx{0ull}, 
-    callstack{}, 
-    stack{}, 
-    program(program), 
+ByteCodeVM::ByteCodeVM(const Program& program) :
+    idx{0ull},
+    function_stack_base{0ull},
+    return_value{0ull},
+    stack{},
+    program(program),
     debug{true},
     ffiFunctions{},
     ffiArgs{} {}
diff --git a/src/executer/ByteCode.h b/src/executer/ByteCode.h
@@ -114,41 +114,14 @@ enum class ProgramState {
     Finished
 };
 
-// TODO: We don't need any of this, we can just use the stack
-// https://stackoverflow.com/questions/1395591/what-exactly-is-the-base-pointer-and-stack-pointer-to-what-do-they-point
-// https://www.cs.virginia.edu/~evans/cs216/guides/x86.html
-
-// We just need to introduce registers
-// esp: Current stack top
-// ebp: Stack base pointer, is set to esp in the beginning of a function
-// eax: Return value of the function
-
-// On function call:
-// - Push the parameters in reverse order
-// - Push return address (done by the CALL instruction)
-// - Push the current ebp
-// - Set ebp to esp
-// - Push locals
-// Within fucnction:
-// - Use ebp + n to access parameters
-// - Use ebp + n to access locals
-// On function return:
-// - Put result into eax register
-// - Pop locals from the stack
-// - Pop ebp from the stack, put into ebp register
-// - Pop return address from the stack, and jump (done by the RET instruction)
-// - Pop parameters from the stack (done by the caller, not the callee)
-// - Get result from eax register
-
-struct StackFrame {
-    std::vector<word_t> locals;
-    word_t return_address;
-};
-
 class ByteCodeVM {
     private:
         word_t idx;
-        std::vector<StackFrame> callstack;
+
+        // Registers to manage stack-based function calls, parameters, locals, and return values
+        word_t function_stack_base;
+        word_t return_value;
+
         Stack stack;
         std::vector<word_t> heap;
         Program program;
diff --git a/src/executer/Stack.cpp b/src/executer/Stack.cpp
@@ -38,4 +38,32 @@ bool Stack::empty() const {
 size_t Stack::size() const {
     return impl.size();
 }
+
+word_t Stack::get(size_t index) const {
+    if (index >= impl.size()) {
+        throwConstraintViolated("Stack::get index out of bounds");
+    }
+    return impl[index];
+}
+
+void Stack::set(size_t index, word_t value) {
+    if (index >= impl.size()) {
+        throwConstraintViolated("Stack::set index out of bounds");
+    }
+    impl[index] = value;
+}
+
+word_t& Stack::operator[](size_t index) {
+    if (index >= impl.size()) {
+        throwConstraintViolated("Stack::operator[] index out of bounds");
+    }
+    return impl[index];
+}
+
+const word_t& Stack::operator[](size_t index) const {
+    if (index >= impl.size()) {
+        throwConstraintViolated("Stack::operator[] const index out of bounds");
+    }
+    return impl[index];
+}
 } // namespace executor
diff --git a/src/executer/Stack.h b/src/executer/Stack.h
@@ -22,6 +22,12 @@ class Stack {
 
     word_t lookback(size_t n) const;
 
+    // Indexed access for register-based VM
+    word_t get(size_t index) const;
+    void set(size_t index, word_t value);
+    word_t& operator[](size_t index);
+    const word_t& operator[](size_t index) const;
+
     using const_iterator = std::vector<word_t>::const_iterator;
 
     auto begin() const {

Original file line number	Diff line number	Diff line change
`@@ -29,6 +29,7 @@ void ByteCodeEmitter::run() {`
`29`	`29`	`for(const auto& param : fn.second->getHead()->getParameters()) {`
`30`	`30`	`localNames.push_back(param->getName());`
`31`	`31`	`}`
	`32`	`+ num_params = localNames.size(); // Remember parameter count`
`32`	`33`	`function_idxs[fn.first] = code().size();`
`33`	`34`	`process(fn.second->getBody(), false);`
`34`	`35`	`}`
`@@ -206,7 +207,9 @@ void ByteCodeEmitter::process(const std::shared_ptr<AST::Node>& node, bool hasCo`
`206`	`207`	`if (ret->getExpr()) {`
`207`	`208`	`process(ret->getExpr(), true);`
`208`	`209`	`}`
`209`		`- code().push_back(executor::Instruction(executor::Op::RET));`
	`210`	`+ // RET arg1=num_params, arg2=num_locals`
	`211`	`+ size_t num_locals = localNames.size() - num_params;`
	`212`	`+ code().push_back(executor::Instruction(executor::Op::RET, num_params, num_locals));`
`210`	`213`	`break;`
`211`	`214`	`}`
`212`	`215`	`case AST::NodeType::Assign: {`