Add tail-call-optimization

2 files changed, 107 insertions, 2 deletions

diff --git a/src/compiler.c b/src/compiler.c
index 11720d5..c2b48b3 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -112,7 +112,8 @@ typedef enum OpCode {
     OP_LDLADDR,  // ldladdr rx, va
     OP_LDSTR,    // ldstr   rx, sa ; Stores the address of the string sa into rx
     // Functions.
-    OP_CALL,     // call fx        ; Bumps the stack pointer by cx
+    OP_CALL,     // call fx        ; Call the function fx
+    OP_RECUR,    // recur          ; Jump to the beginning of the function.
     OP_RET,      // ret            ; Returns from current function
     OP_RESERVE,  // reserve cx     ; Increments the stack pointer by cx bytes
     OP_POP,      // pop     rx     ; Pops the last value of the stack into rx.
@@ -241,6 +242,7 @@ Str op_str[] = {
     [OP_PUTRETI] = cstr("PUTRETI "),
     // Functions.
     [OP_CALL] = cstr("CALL    "),
+    [OP_RECUR] = cstr("RECUR   "),
     [OP_RET] = cstr("RET     "),
     [OP_RESERVE] = cstr("RESERVE "),
     [OP_POP] = cstr("POP     "),
@@ -900,6 +902,44 @@ compile_while(Compiler *compiler, Chunk *chunk, Node *node) {
 }
 
 CompResult
+compile_tail_call(Compiler *compiler, Chunk *chunk, Node *node) {
+    // Update the local parameters.
+    for (sz i = 0; i < array_size(node->elements); i++) {
+        Node *expr = node->elements[i];
+        CompResult result = compile_expr(compiler, chunk, expr);
+        switch (result.type) {
+            case COMP_CONST: {
+                emit_op(OP_STLVARI, i, result.idx, 0, node, chunk);
+            } break;
+            case COMP_REG: {
+                if (str_eq(expr->type, cstr("str"))) {
+                    sz var_addr = chunk->reg_idx++;
+                    sz str_addr = result.idx;
+                    emit_op(OP_LDLADDR, var_addr, i, 0, node, chunk);
+                    emit_fat_copy(chunk, node, var_addr, str_addr);
+                } else {
+                    emit_op(OP_STLVAR, i, result.idx, 0, node, chunk);
+                }
+            } break;
+            case COMP_STRING: {
+                sz var_addr = chunk->reg_idx++;
+                sz str_addr = chunk->reg_idx++;
+                emit_op(OP_LDLADDR, var_addr, i, 0, node, chunk);
+                emit_op(OP_LDSTR, str_addr, result.idx, 0, node, chunk);
+                emit_fat_copy(chunk, node, var_addr, str_addr);
+            } break;
+            default: {
+                emit_compile_err(compiler, chunk, node);
+                return (CompResult){.type = COMP_ERR};
+            } break;
+        }
+    }
+
+    emit_op(OP_RECUR, 0, 0, 0, node, chunk);
+    return (CompResult){.type = COMP_NIL};
+}
+
+CompResult
 compile_funcall(Compiler *compiler, Chunk *chunk, Node *node) {
     Str name = node->value.str;
 
@@ -977,6 +1017,67 @@ compile_funcall(Compiler *compiler, Chunk *chunk, Node *node) {
     }
     Function fun = map->val;
 
+    // Check for tail recursive opportunities.
+    if (str_eq(fun.name, node->unique_name) &&
+        str_eq(chunk->name, node->unique_name)) {
+        Node *parent = node->parent;
+        Node *current = node;
+        bool tail_recursive = true;
+        while (parent != NULL) {
+            switch (parent->kind) {
+                case NODE_BLOCK: {
+                    sz idx = array_size(parent->statements) - 1;
+                    if (parent->statements[idx] != node) {
+                        tail_recursive = false;
+                        break;
+                    }
+                } break;
+                case NODE_WHILE: {
+                    if (current == parent->loop.cond) {
+                        tail_recursive = false;
+                        break;
+                    }
+                } break;
+                case NODE_IF: {
+                    if (current == parent->ifelse.cond) {
+                        tail_recursive = false;
+                        break;
+                    }
+                } break;
+                case NODE_FUN: {
+                    sz idx = array_size(parent->func.body->statements) - 1;
+                    if (parent->func.body->statements[idx] != current) {
+                        tail_recursive = false;
+                        break;
+                    }
+                    break;
+                } break;
+                case NODE_MATCH: {
+                    if (current == parent->match.expr) {
+                        tail_recursive = false;
+                        break;
+                    }
+                } break;
+                case NODE_COND: break;
+                case NODE_CASE_COND: {
+                    if (current == parent->case_entry.cond) {
+                        tail_recursive = false;
+                        break;
+                    }
+                } break;
+                default: {
+                    tail_recursive = false;
+                    break;
+                } break;
+            }
+            parent = parent->parent;
+            current = current->parent;
+        }
+        if (tail_recursive) {
+            return compile_tail_call(compiler, chunk, node);
+        }
+    }
+
     // Reserve space for the return value if needed.
     if (fun.return_arity > 0) {
         // Put the return data into a register
@@ -1655,6 +1756,7 @@ disassemble_instruction(Instruction instruction) {
                     instruction.a, instruction.b);
             break;
         case OP_RET:
+        case OP_RECUR:
         case OP_HALT: println("%s", op_str[instruction.op]); break;
         default: println("Unknown opcode %d", instruction.op); break;
     }
diff --git a/src/vm.c b/src/vm.c
index 0929fb6..0791706 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -5,7 +5,7 @@
 #include "compiler.c"
 
 #define N_CONST 256
-#define STACK_SIZE KB(64)
+#define STACK_SIZE MB(4)
 typedef struct VM {
     Chunk *main;
     Chunk *chunk;
@@ -382,6 +382,9 @@ vm_run(VM *vm) {
                 p[3] = old_fp;
                 vm->sp += sizeof(ptrsize) * 4;
             } break;
+            case OP_RECUR: {
+                vm->ip = vm->chunk->code;
+            } break;
             case OP_RET: {
                 u64 *p = (u64 *)vm->sp;
                 ptrsize chunk_addr = p[-4];