1 files changed, 41 insertions, 26 deletions

diff --git a/src/compiler.c b/src/compiler.c
index 5bddcb6..f46d00d 100644
--- a/src/compiler.c
+++ b/src/compiler.c
@@ -504,34 +504,49 @@ compile_if(Chunk *chunk, Node *node) {
         }
     }
 
-    // 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;
+    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);
+
+        // 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_else, chunk);
+                } break;
+                case COMP_REG: {
+                    EMIT_OP(OP_MOV64, reg_dst, else_expr.idx, 0,
+                            node->cond_else, 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;
+    } 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);
+            EMIT_OP(OP_LD64K, reg_dst, zero, 0, node, chunk);
+        } else {
+            sz const_a = add_constant(chunk, end_expr - jump_a);
+            chunk->code[jump_a].a = const_a;
         }
     }
-    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?