Change jumps to be label based from backpatching

3 files changed, 117 insertions, 85 deletions

diff --git a/src/compiler.c b/src/compiler.c
index ea9234a..4c6a946 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -25,7 +25,7 @@ typedef struct Instruction {
 typedef union Constant {
     s64 i;
     u64 u;
-    double f;
+    f64 f;
     ptrsize ptr;
 } Constant;
 
@@ -37,6 +37,9 @@ typedef struct LineCol {
 typedef struct Chunk {
     sz id;
     Instruction *code;
+    IntIntMap *labels;      // label -> chunk_index
+    IntIntMap *labels_rev;  // chunk_index -> label
+    sz labels_idx;
 
     // Constant values that fit in 64 bits.
     Constant *constants;
@@ -149,12 +152,11 @@ typedef enum OpCode {
     OP_BITOR,       // bor    rx, ra, rb
     OP_BITNOT,      // bnot   rx, ra
     // Jump instructions.
-    OP_JMPI,   // jmp  cx     ; cx := signed offset
-    OP_JMPFI,  // jmpf cx, ca ; rx := condition, ca := offset
-    OP_JMPTI,  // jmpt cx, ca ; rx := condition, ca := offset
-    OP_JMP,    // jmp  rx     ; rx := signed offset
-    OP_JMPF,   // jmpf rx, ca ; rx := condition, ca  := offset
-    OP_JMPT,   // jmpt rx, ca ; rx := condition, ca  := offset
+    OP_JMP,    // jmp  lx     ; jmp to label lx
+    OP_JMPF,   // jmpf lx, rx ; jmp to label lx if rx is false
+    OP_JMPT,   // jmpt lx, rx ; jmp to label lx if rx is true
+    OP_JMPFI,  // jmpf lx, cx ; jmp to label lx if rx is false
+    OP_JMPTI,  // jmpt lx, cx ; jmp to label lx if rx is true
 } OpCode;
 
 Str op_str[] = {
@@ -241,12 +243,11 @@ Str op_str[] = {
     [OP_BITOR] = cstr("BOR     "),
     [OP_BITNOT] = cstr("BNOT    "),
     // Jump instructions.
-    [OP_JMPI] = cstr("JMPI    "),
-    [OP_JMPFI] = cstr("JMPFI   "),
-    [OP_JMPTI] = cstr("JMPTI   "),
     [OP_JMP] = cstr("JMP     "),
     [OP_JMPF] = cstr("JMPF    "),
     [OP_JMPT] = cstr("JMPT    "),
+    [OP_JMPFI] = cstr("JMPFI   "),
+    [OP_JMPTI] = cstr("JMPTI   "),
 };
 
 typedef enum {
@@ -264,17 +265,21 @@ typedef struct CompResult {
 
 CompResult compile_expr(Chunk *chunk, Node *node);
 
-#define EMIT_OP(OP, DST, A, B, NODE, CHUNK)                                    \
-    do {                                                                       \
-        Instruction inst = (Instruction){                                      \
-            .op = (OP),                                                        \
-            .dst = (DST),                                                      \
-            .a = (A),                                                          \
-            .b = (B),                                                          \
-        };                                                                     \
-        array_push((CHUNK)->code, inst, (CHUNK)->storage);                     \
-        LineCol linecol = (LineCol){.line = (NODE)->line, .col = (NODE)->col}; \
-        array_push((CHUNK)->linecol, linecol, (CHUNK)->storage);               \
+#define EMIT_OP(OP, DST, A, B, NODE, CHUNK)                              \
+    do {                                                                 \
+        Instruction inst = (Instruction){                                \
+            .op = (OP),                                                  \
+            .dst = (DST),                                                \
+            .a = (A),                                                    \
+            .b = (B),                                                    \
+        };                                                               \
+        array_push((CHUNK)->code, inst, (CHUNK)->storage);               \
+        LineCol linecol = (LineCol){0};                                  \
+        if (NODE) {                                                      \
+            Node *_node = (NODE);                                        \
+            linecol = (LineCol){.line = _node->line, .col = _node->col}; \
+        }                                                                \
+        array_push((CHUNK)->linecol, linecol, (CHUNK)->storage);         \
     } while (0)
 
 CompResult
@@ -476,7 +481,6 @@ compile_if(Chunk *chunk, Node *node) {
             return (CompResult){.type = COMP_ERR};
         } break;
     }
-    sz jump_a = array_size(chunk->code);
     sz reg_dst = 255;
     bool has_value = !str_eq(node->type, cstr("nil"));
     if (has_value) {
@@ -484,7 +488,8 @@ compile_if(Chunk *chunk, Node *node) {
     }
 
     // Jump to the `false` branch.
-    EMIT_OP(jmpop, cond.idx, 0xff, 0, node->cond_if, chunk);
+    sz lab0 = chunk->labels_idx++;
+    EMIT_OP(jmpop, lab0, cond.idx, 0, node->cond_if, chunk);
 
     // Condition is true.
     CompResult then_expr = compile_expr(chunk, node->cond_expr);
@@ -506,8 +511,9 @@ compile_if(Chunk *chunk, Node *node) {
 
     if (node->cond_else) {
         // Jump to the end of the expression.
-        sz jump_b = array_size(chunk->code);
-        EMIT_OP(OP_JMPI, 0xff, 0, 0, node->cond_else, chunk);
+        sz pos0 = array_size(chunk->code);
+        sz lab1 = chunk->labels_idx++;
+        EMIT_OP(OP_JMP, lab1, 0, 0, node->cond_else, chunk);
 
         // Else expression.
         CompResult else_expr = compile_expr(chunk, node->cond_else);
@@ -526,25 +532,33 @@ compile_if(Chunk *chunk, Node *node) {
                 } break;
             }
         }
-        sz end_expr = array_size(chunk->code);
+        sz pos1 = array_size(chunk->code);
 
-        // Backpatch jumps.
-        sz const_a = add_constant(chunk, jump_b + 1 - jump_a);
-        sz const_b = add_constant(chunk, end_expr - jump_b);
-        chunk->code[jump_a].a = const_a;
-        chunk->code[jump_b].dst = const_b;
+        // Update labels.
+        intintmap_insert(&chunk->labels, lab0, pos0 + 1, chunk->storage);
+        intintmap_insert(&chunk->labels, lab1, pos1, chunk->storage);
+        intintmap_insert(&chunk->labels_rev, pos0 + 1, lab0, chunk->storage);
+        intintmap_insert(&chunk->labels_rev, pos1, lab1, chunk->storage);
     } else {
-        sz end_expr = array_size(chunk->code);
         if (has_value) {
-            sz const_a = add_constant(chunk, end_expr + 1 - jump_a);
-            chunk->code[jump_a].a = const_a;
             sz zero = add_constant(chunk, 0);
-            sz end = add_constant(chunk, 2);
-            EMIT_OP(OP_JMPI, end, 0, 0, node, chunk);
+            sz lab1 = chunk->labels_idx++;
+            EMIT_OP(OP_JMP, lab1, 0, 0, node, chunk);
+            sz pos0 = array_size(chunk->code);
             EMIT_OP(OP_LD64K, reg_dst, zero, 0, node, chunk);
+            sz pos1 = array_size(chunk->code);
+
+            // Update labels.
+            intintmap_insert(&chunk->labels, lab0, pos0, chunk->storage);
+            intintmap_insert(&chunk->labels, lab1, pos1, chunk->storage);
+            intintmap_insert(&chunk->labels_rev, pos0, lab0, chunk->storage);
+            intintmap_insert(&chunk->labels_rev, pos1, lab1, chunk->storage);
         } else {
-            sz const_a = add_constant(chunk, end_expr - jump_a);
-            chunk->code[jump_a].a = const_a;
+            sz pos0 = array_size(chunk->code);
+
+            // Update labels.
+            intintmap_insert(&chunk->labels, lab0, pos0, chunk->storage);
+            intintmap_insert(&chunk->labels_rev, pos0, lab0, chunk->storage);
         }
     }
     // TODO: does it has an else or not? Moreover, should we enforce on the
@@ -560,7 +574,7 @@ compile_if(Chunk *chunk, Node *node) {
 
 CompResult
 compile_while(Chunk *chunk, Node *node) {
-    sz start = array_size(chunk->code);
+    sz pos1 = array_size(chunk->code);
     CompResult cond = compile_expr(chunk, node->while_cond);
     OpCode jmpop;
     switch (cond.type) {
@@ -574,18 +588,22 @@ compile_while(Chunk *chunk, Node *node) {
             return (CompResult){.type = COMP_ERR};
         } break;
     }
-    sz jump_a = array_size(chunk->code);
 
     // Jump to the `end of the loop` branch.
-    EMIT_OP(jmpop, cond.idx, 0xff, 0, node->while_cond, chunk);
+    sz lab0 = chunk->labels_idx++;
+    EMIT_OP(jmpop, lab0, cond.idx, 0, node->while_cond, chunk);
 
     // Condition is true.
-    CompResult then_expr = compile_expr(chunk, node->while_expr);
-    sz end_expr = array_size(chunk->code);
-    sz loopback = add_constant(chunk, (start - end_expr));
-    EMIT_OP(OP_JMPI, loopback, 0, 0, node, chunk);
-    sz const_a = add_constant(chunk, end_expr - jump_a + 1);
-    chunk->code[jump_a].a = const_a;
+    compile_expr(chunk, node->while_expr);
+    sz pos0 = array_size(chunk->code);
+    sz lab1 = chunk->labels_idx++;
+    EMIT_OP(OP_JMP, lab1, 0, 0, node, chunk);
+
+    // Update labels.
+    intintmap_insert(&chunk->labels, lab0, pos0 + 1, chunk->storage);
+    intintmap_insert(&chunk->labels, lab1, pos1, chunk->storage);
+    intintmap_insert(&chunk->labels_rev, pos0 + 1, lab0, chunk->storage);
+    intintmap_insert(&chunk->labels_rev, pos1, lab1, chunk->storage);
 
     // Return.
     return (CompResult){.type = COMP_NIL};
@@ -747,12 +765,15 @@ disassemble_instruction(Instruction instruction) {
             println("%s r%d, r%d", op_str[instruction.op], instruction.dst,
                     instruction.a, instruction.b);
             break;
+        case OP_JMPF:
+        case OP_JMPT:
+            println("%s l%d, r%d", op_str[instruction.op], instruction.dst,
+                    instruction.a, instruction.b);
+            break;
         case OP_LD8K:
         case OP_LD16K:
         case OP_LD32K:
         case OP_LD64K:
-        case OP_JMPF:
-        case OP_JMPT:
             println("%s r%d, c%d", op_str[instruction.op], instruction.dst,
                     instruction.a, instruction.b);
             break;
@@ -844,17 +865,13 @@ disassemble_instruction(Instruction instruction) {
             println("%s r%d, r%d", op_str[instruction.op], instruction.dst,
                     instruction.a, instruction.b);
             break;
-        case OP_JMPI:
-            println("%s c%d", op_str[instruction.op], instruction.dst,
-                    instruction.a, instruction.b);
-            break;
         case OP_JMP:
-            println("%s r%d", op_str[instruction.op], instruction.dst,
+            println("%s l%d", op_str[instruction.op], instruction.dst,
                     instruction.a, instruction.b);
             break;
         case OP_JMPFI:
         case OP_JMPTI:
-            println("%s c%d, c%d", op_str[instruction.op], instruction.dst,
+            println("%s l%d, c%d", op_str[instruction.op], instruction.dst,
                     instruction.a, instruction.b);
             break;
         case OP_HALT: println("%s", op_str[instruction.op]); break;
@@ -864,26 +881,39 @@ disassemble_instruction(Instruction instruction) {
 
 void
 disassemble_chunk(Chunk chunk) {
-    println("%s: ============== code ==============", chunk.file_name);
+    println("%s: ================== code ==================", chunk.file_name);
+    println("%s: LINE:COL  INUM LABELS OP       OPERANDS   ", chunk.file_name);
+    println("%s: ------------------------------------------", chunk.file_name);
     for (sz i = 0; i < array_size(chunk.code); i++) {
-        print("%s: %x{4}:  ", chunk.file_name, i);
+        print("%s:", chunk.file_name);
+        printf(" %.4ld:%.4ld %.4lx ", chunk.linecol[i].line,
+               chunk.linecol[i].col, i);
+        IntIntMap *label = intintmap_lookup(&chunk.labels_rev, i);
+        if (label) {
+            printf(".L%.2ld   ", label->val);
+        } else {
+            printf("  %2s   ", "");
+        }
         disassemble_instruction(chunk.code[i]);
     }
     if (array_size(chunk.constants) > 0) {
-        println("%s: ============ constants ===========", chunk.file_name);
+        println("%s: ================ constants ===============",
+                chunk.file_name);
         for (sz i = 0; i < array_size(chunk.constants); i++) {
             println("%s: %x{2}:  %x{8}", chunk.file_name, i,
                     chunk.constants[i]);
         }
     }
     if (array_size(chunk.strings) > 0) {
-        println("%s: ============= strings ============", chunk.file_name);
+        println("%s: ================= strings ================",
+                chunk.file_name);
         for (sz i = 0; i < array_size(chunk.strings); i++) {
             println("%s: %x{2}:  %s", chunk.file_name, i, chunk.strings[i]);
         }
     }
     if (array_size(chunk.vars) > 0) {
-        println("%s: ============ variables ===========", chunk.file_name);
+        println("%s: ================ variables ===============",
+                chunk.file_name);
         for (sz i = 0; i < array_size(chunk.vars); i++) {
             println("%s: %x{2}: [%x{4}:%x{4}] %s: %s", chunk.file_name, i,
                     chunk.vars[i].offset,
@@ -891,6 +921,7 @@ disassemble_chunk(Chunk chunk) {
                     chunk.vars[i].name, chunk.vars[i].type);
         }
     }
+    println("%s: ==========================================", chunk.file_name);
     println("n_regs: %d, n_vars: %d, n_strings: %d, n_consts: %d",
             chunk.reg_idx, array_size(chunk.vars), chunk.str_idx,
             chunk.const_idx);
diff --git a/src/main.c b/src/main.c
index 0e453c6..5e5775e 100644
--- a/src/main.c
+++ b/src/main.c
@@ -181,6 +181,7 @@ process_file(Str path) {
             exit(EXIT_FAILURE);
         }
     }
+    // Make sure the last result is on r0.
     sz res_reg = 0;
     switch (res.type) {
         case COMP_CONST: {
@@ -194,9 +195,7 @@ process_file(Str path) {
         } break;
         default: break;
     }
-    // After we are done move the last result to r0 for printing.
-    Instruction halt = (Instruction){.op = OP_HALT, .dst = res_reg};
-    array_push(chunk.code, halt, &bytecode_arena);
+    EMIT_OP(OP_HALT, res_reg, 0, 0, NULL, &chunk);
 
     if (chunk.const_idx >= 256) {
         eprintln("too many constants on chunk %s", chunk.id);
diff --git a/src/vm.c b/src/vm.c
index 205c15a..52f5476 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -155,40 +155,41 @@ vm_run(VM *vm) {
                 s64 *stack = (s64 *)&vm->stack[var.offset];
                 *stack = vm->chunk->constants[src].i;
             } break;
-            case OP_JMPI: {
-                sz offset = vm->chunk->constants[instruction.dst].i;
-                vm->ip += offset - 1;
+            case OP_JMP: {
+                u8 dst = instruction.dst;
+                sz pos = intintmap_lookup(&vm->chunk->labels, dst)->val;
+                vm->ip = vm->chunk->code + pos;
             } break;
             case OP_JMPFI: {
-                bool cond = vm->chunk->constants[instruction.dst].i;
-                sz offset = vm->chunk->constants[instruction.a].i;
+                u8 dst = instruction.dst;
+                sz pos = intintmap_lookup(&vm->chunk->labels, dst)->val;
+                bool cond = vm->chunk->constants[instruction.a].i;
                 if (!cond) {
-                    vm->ip += offset - 1;
+                    vm->ip = vm->chunk->code + pos;
                 }
             } break;
             case OP_JMPTI: {
-                bool cond = vm->chunk->constants[instruction.dst].i;
-                sz offset = vm->chunk->constants[instruction.a].i;
+                u8 dst = instruction.dst;
+                sz pos = intintmap_lookup(&vm->chunk->labels, dst)->val;
+                bool cond = vm->chunk->constants[instruction.a].i;
                 if (cond) {
-                    vm->ip += offset - 1;
+                    vm->ip = vm->chunk->code + pos;
                 }
             } break;
-            case OP_JMP: {
-                sz offset = vm->chunk->constants[instruction.dst].i;
-                vm->ip += offset - 1;
-            } break;
             case OP_JMPF: {
-                bool cond = vm->regs[instruction.dst].i;
-                sz offset = vm->chunk->constants[instruction.a].i;
+                u8 dst = instruction.dst;
+                sz pos = intintmap_lookup(&vm->chunk->labels, dst)->val;
+                bool cond = vm->regs[instruction.a].i;
                 if (!cond) {
-                    vm->ip += offset - 1;
+                    vm->ip = vm->chunk->code + pos;
                 }
             } break;
             case OP_JMPT: {
-                bool cond = vm->regs[instruction.dst].i;
-                sz offset = vm->chunk->constants[instruction.a].i;
+                u8 dst = instruction.dst;
+                sz pos = intintmap_lookup(&vm->chunk->labels, dst)->val;
+                bool cond = vm->regs[instruction.a].i;
                 if (cond) {
-                    vm->ip += offset - 1;
+                    vm->ip = vm->chunk->code + pos;
                 }
             } break;
             case OP_MOV64: {
@@ -212,9 +213,10 @@ vm_run(VM *vm) {
                 vm->regs[dst].i = vm->regs[src].i & 0xFF;
             } break;
             case OP_HALT: {
-                println("VM HALT (int)   -> %d", vm->regs[instruction.dst]);
-                println("VM HALT (float) -> %f", vm->regs[instruction.dst]);
-                println("VM HALT (hex)   -> %x", vm->regs[instruction.dst]);
+                Constant result = vm->regs[instruction.dst];
+                printf("VM HALT (int)   -> %lld\n", result.i);
+                printf("VM HALT (float) -> %.10e\n", result.f);
+                printf("VM HALT (hex)   -> %llx\n", (u64)result.i);
                 return;
             }
             default: {