Cleanup IF jump code in compiler

3 files changed, 41 insertions, 24 deletions

diff --git a/src/bytecode/compiler.h b/src/bytecode/compiler.h
index 559ad1c..c01f745 100755
--- a/src/bytecode/compiler.h
+++ b/src/bytecode/compiler.h
@@ -45,6 +45,22 @@ emit_constant(Chunk *chunk, Token tok, Object obj) {
     }
 }
 
+size_t
+emit_jump(Chunk *chunk, Token tok, Ops op) {
+    add_code(chunk, op, tok.line, tok.column);
+    add_code(chunk, 0xFF, tok.line, tok.column);
+    add_code(chunk, 0xFF, tok.line, tok.column);
+    return array_size(chunk->code) - 2;
+}
+
+void
+patch_jump(Chunk *chunk, size_t offset) {
+  size_t jump = array_size(chunk->code) - offset - 2;
+  assert(jump <= UINT16_MAX && "error: jump is too long");
+  chunk->code[offset] = (jump >> 8) & 0xFF;
+  chunk->code[offset + 1] = (jump) & 0xFF;
+}
+
 void
 parse_fixnum(Chunk *chunk, Token tok) {
     ssize_t num = 0;
@@ -255,34 +271,23 @@ compile_if_op(Chunk *chunk, Visitor *vs, Token start) {
 
     // Condition.
     parse_tree(chunk, vs);
-    emit_constant(chunk, tok, FIXNUM_VAL(0xFF));
-    size_t pos = array_size(chunk->code);
-    add_code(chunk, OP_JUMP_IF_FALSE, start.line, start.column);
+    size_t jmp_false = emit_jump(chunk, start, OP_JUMP_IF_FALSE);
 
-    // Block true.
+    // True expression.
     parse_tree(chunk, vs);
 
-    // No else.
+    // No second expression.
     if (peek_token(vs).type == TOKEN_RPAREN) {
-        size_t false_block_start = array_size(chunk->code) - pos - 1;
-        size_t num_idx = add_constant(chunk, FIXNUM_VAL(false_block_start));
-        chunk->code[pos - 1] = num_idx;
+        patch_jump(chunk, jmp_false);
         next_token(vs);
         return;
     }
 
-    // Else.
-    emit_constant(chunk, tok, FIXNUM_VAL(0xFF));
-    size_t else_pos = array_size(chunk->code);
-    add_code(chunk, OP_JUMP, start.line, start.column);
-
-    size_t false_block_start = array_size(chunk->code) - pos - 1;
-    size_t num_idx = add_constant(chunk, FIXNUM_VAL(false_block_start));
-    chunk->code[pos - 1] = num_idx;
+    // False expression.
+    size_t jmp_end = emit_jump(chunk, start, OP_JUMP);
+    patch_jump(chunk, jmp_false);
     parse_tree(chunk, vs);
-    false_block_start = array_size(chunk->code) - else_pos - 1;
-    num_idx = add_constant(chunk, FIXNUM_VAL(false_block_start));
-    chunk->code[else_pos - 1] = num_idx;
+    patch_jump(chunk, jmp_end);
 
     if (peek_token(vs).type != TOKEN_RPAREN) {
         error_push((Error){
diff --git a/src/bytecode/debug.h b/src/bytecode/debug.h
index 674d469..bc736f9 100755
--- a/src/bytecode/debug.h
+++ b/src/bytecode/debug.h
@@ -71,11 +71,19 @@ disassemble_instruction(Chunk *chunk, size_t offset) {
     switch (instruction) {
         case OP_CONSTANT: {
             u8 constant = chunk->code[offset + 1];
-            printf("%-16s %4d -> ", "OP_CONSTANT", constant);
+            printf("%-16s %4d -> ", ops_str[instruction], constant);
             display(chunk->constants[constant]);
             printf("\n");
             return offset + 2;
         } break;
+        case OP_JUMP:
+        case OP_JUMP_IF_FALSE: {
+            u16 a = chunk->code[offset + 1];
+            u16 b = chunk->code[offset + 2];
+            u16 jmp = (a << 8) | b;
+            printf("%-16s %4d\n", ops_str[instruction], jmp);
+            return offset + 3;
+        } break;
         default: {
             printf("%s\n", ops_str[instruction]);
             return offset + 1;
diff --git a/src/bytecode/vm.h b/src/bytecode/vm.h
index b24a5fa..ba33077 100755
--- a/src/bytecode/vm.h
+++ b/src/bytecode/vm.h
@@ -197,13 +197,17 @@ vm_interpret(VM *vm, Chunk *chunk) {
             case OP_LESS_EQUAL: { FIXNUM_COMPARE_OP(<=); } break;
             case OP_GREATER_EQUAL: { FIXNUM_COMPARE_OP(>=); } break;
             case OP_JUMP: {
-                ssize_t off = AS_FIXNUM(array_pop(vm->stack));
-                vm->pc += off;
+                u16 a = *vm->pc++;
+                u16 b = *vm->pc++;
+                u16 offset = (a << 8) | b;
+                vm->pc += offset;
             } break;
             case OP_JUMP_IF_FALSE: {
-                ssize_t off = AS_FIXNUM(array_pop(vm->stack));
+                u16 a = *vm->pc++;
+                u16 b = *vm->pc++;
+                u16 offset = (a << 8) | b;
                 if (IS_FALSE(array_pop(vm->stack))) {
-                    vm->pc += off;
+                    vm->pc += offset;
                 }
             } break;
             case OP_DISPLAY: {