Add if/else expressions

2 files changed, 185 insertions, 11 deletions

diff --git a/src/compiler.c b/src/compiler.c
index 84cc99a..5bddcb6 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -148,6 +148,13 @@ typedef enum OpCode {
     OP_BITAND,      // band   rx, ra, rb
     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
 } OpCode;
 
 Str op_str[] = {
@@ -233,6 +240,13 @@ Str op_str[] = {
     [OP_BITAND] = cstr("BAND    "),
     [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    "),
 };
 
 typedef enum {
@@ -434,9 +448,105 @@ compile_unary(Chunk *chunk, Node *node) {
     return (CompResult){.type = COMP_REG, .idx = reg_dst};
 }
 
+sz
+add_constant(Chunk *chunk, sz value) {
+    IntIntMap *map = intintmap_lookup(&chunk->intmap, value);
+    // Make sure we don't have duplicated constants.
+    if (!map) {
+        map = intintmap_insert(&chunk->intmap, value, chunk->const_idx++,
+                               chunk->storage);
+        Constant c = (Constant){.i = value};
+        array_push(chunk->constants, c, chunk->storage);
+    }
+    return map->val;
+}
+
+CompResult
+compile_if(Chunk *chunk, Node *node) {
+    CompResult cond = compile_expr(chunk, node->cond_if);
+    OpCode jmpop;
+    switch (cond.type) {
+        case COMP_CONST: {
+            jmpop = OP_JMPFI;
+        } break;
+        case COMP_REG: {
+            jmpop = OP_JMPF;
+        } break;
+        default: {
+            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) {
+        reg_dst = chunk->reg_idx++;
+    }
+
+    // Jump to the `false` branch.
+    EMIT_OP(jmpop, cond.idx, 0xff, 0, node->cond_if, chunk);
+
+    // Condition is true.
+    CompResult then_expr = compile_expr(chunk, node->cond_expr);
+    if (has_value) {
+        switch (then_expr.type) {
+            case COMP_CONST: {
+                EMIT_OP(OP_LD64K, reg_dst, then_expr.idx, 0, node->cond_if,
+                        chunk);
+            } break;
+            case COMP_REG: {
+                EMIT_OP(OP_MOV64, reg_dst, then_expr.idx, 0, node->cond_if,
+                        chunk);
+            } break;
+            default: {
+                return (CompResult){.type = COMP_ERR};
+            } break;
+        }
+    }
+
+    // Jump to the end of the expression.
+    sz jump_b = array_size(chunk->code);
+    EMIT_OP(OP_JMPI, 0xff, 0, 0, node->cond_if, chunk);
+
+    // Else expression.
+    CompResult else_expr = compile_expr(chunk, node->cond_else);
+    if (has_value) {
+        switch (else_expr.type) {
+            case COMP_CONST: {
+                EMIT_OP(OP_LD64K, reg_dst, else_expr.idx, 0, node->cond_if,
+                        chunk);
+            } break;
+            case COMP_REG: {
+                EMIT_OP(OP_MOV64, reg_dst, else_expr.idx, 0, node->cond_if,
+                        chunk);
+            } break;
+            default: {
+                return (CompResult){.type = COMP_ERR};
+            } break;
+        }
+    }
+    sz end_expr = 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;
+    // TODO: does it has an else or not? Moreover, should we enforce on the
+    // semantic level that if the `if` expression returns a value we must add an
+    // else?
+
+    // Return.
+    if (has_value) {
+        return (CompResult){.type = COMP_REG, .idx = reg_dst};
+    }
+    return (CompResult){.type = COMP_NIL};
+}
+
 CompResult
 compile_expr(Chunk *chunk, Node *node) {
     switch (node->kind) {
+        case NODE_IF: return compile_if(chunk, node);
         // Logic.
         // case NODE_XOR:
         case NODE_BITNOT:
@@ -466,17 +576,10 @@ compile_expr(Chunk *chunk, Node *node) {
         case NODE_NUM_UINT:
         case NODE_NUM_INT: {
             sz value = node->value.i;
-            // Make sure we don't have duplicated constants.
-            IntIntMap *map = intintmap_lookup(&chunk->intmap, value);
-            if (!map) {
-                map = intintmap_insert(&chunk->intmap, value,
-                                       chunk->const_idx++, chunk->storage);
-                Constant c = (Constant){.i = node->value.i};
-                array_push(chunk->constants, c, chunk->storage);
-            }
+            sz const_idx = add_constant(chunk, value);
             return (CompResult){
                 .type = COMP_CONST,
-                .idx = map->val,
+                .idx = const_idx,
             };
         } break;
         case NODE_STRING: {
@@ -577,6 +680,8 @@ disassemble_instruction(Instruction instruction) {
         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;
@@ -668,6 +773,19 @@ 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,
+                    instruction.a, instruction.b);
+            break;
+        case OP_JMPFI:
+        case OP_JMPTI:
+            println("%s c%d, c%d", op_str[instruction.op], instruction.dst,
+                    instruction.a, instruction.b);
+            break;
         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 4e15bee..205c15a 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -137,7 +137,6 @@ vm_run(VM *vm) {
             case OP_LDGVAR: {
                 u8 dst = instruction.dst;
                 u8 src = instruction.a;
-                println("dst: %d src: %d", dst, src);
                 Variable var = vm->chunk->vars[src];
                 s64 *stack = (s64 *)&vm->stack[var.offset];
                 vm->regs[dst].i = *stack;
@@ -156,6 +155,62 @@ 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;
+            } break;
+            case OP_JMPFI: {
+                bool cond = vm->chunk->constants[instruction.dst].i;
+                sz offset = vm->chunk->constants[instruction.a].i;
+                if (!cond) {
+                    vm->ip += offset - 1;
+                }
+            } break;
+            case OP_JMPTI: {
+                bool cond = vm->chunk->constants[instruction.dst].i;
+                sz offset = vm->chunk->constants[instruction.a].i;
+                if (cond) {
+                    vm->ip += offset - 1;
+                }
+            } 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;
+                if (!cond) {
+                    vm->ip += offset - 1;
+                }
+            } break;
+            case OP_JMPT: {
+                bool cond = vm->regs[instruction.dst].i;
+                sz offset = vm->chunk->constants[instruction.a].i;
+                if (cond) {
+                    vm->ip += offset - 1;
+                }
+            } break;
+            case OP_MOV64: {
+                u8 dst = instruction.dst;
+                u8 src = instruction.a;
+                vm->regs[dst] = vm->regs[src];
+            } break;
+            case OP_MOV32: {
+                u8 dst = instruction.dst;
+                u8 src = instruction.a;
+                vm->regs[dst].i = vm->regs[src].i & 0xFFFFFFFF;
+            } break;
+            case OP_MOV16: {
+                u8 dst = instruction.dst;
+                u8 src = instruction.a;
+                vm->regs[dst].i = vm->regs[src].i & 0xFFFF;
+            } break;
+            case OP_MOV8: {
+                u8 dst = instruction.dst;
+                u8 src = instruction.a;
+                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]);
@@ -163,7 +218,8 @@ vm_run(VM *vm) {
                 return;
             }
             default: {
-                eprintln("unimplemented OP code: %d", instruction.op);
+                // eprintln("unimplemented OP code: %d", instruction.op);
+                eprintln("unimplemented OP code: %s", op_str[instruction.op]);
                 return;
             }
         }