From 3143e8c203de6bcdfcc6226a2836e46a07af4980 Mon Sep 17 00:00:00 2001
From: Jeff <dev@jeffa.io>
Date: Sat, 16 Nov 2024 05:40:31 -0500
Subject: [PATCH] Fix control flow and pass all tests

---
 dust-lang/src/compiler.rs       | 17 ++++++++++++-----
 dust-lang/tests/control_flow.rs | 28 ++++++++++++++++++----------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/dust-lang/src/compiler.rs b/dust-lang/src/compiler.rs
index 5ef781b..04c5f88 100644
--- a/dust-lang/src/compiler.rs
+++ b/dust-lang/src/compiler.rs
@@ -270,11 +270,17 @@ impl<'src> Compiler<'src> {
         Some(n_operations)
     }
 
-    fn get_last_jumpable_mut(&mut self) -> Option<&mut Instruction> {
+    fn get_last_jumpable_mut_between(
+        &mut self,
+        minimum: usize,
+        maximum: usize,
+    ) -> Option<&mut Instruction> {
         self.chunk
             .instructions_mut()
             .iter_mut()
             .rev()
+            .skip(minimum)
+            .take(maximum)
             .find_map(|(instruction, _)| {
                 if let Operation::LoadBoolean | Operation::LoadConstant = instruction.operation() {
                     Some(instruction)
@@ -1124,7 +1130,7 @@ impl<'src> Compiler<'src> {
         let if_block_type = self.previous_expression_type.clone();
         let if_last_register = self.next_register().saturating_sub(1);
 
-        let has_else_statement = if let Token::Else = self.current_token {
+        if let Token::Else = self.current_token {
             self.advance()?;
 
             if let Token::LeftBrace = self.current_token {
@@ -1158,8 +1164,10 @@ impl<'src> Compiler<'src> {
 
         match else_block_distance {
             0 => {}
-            1 if !has_else_statement => {
-                if let Some(skippable) = self.get_last_jumpable_mut() {
+            1 => {
+                if let Some(skippable) =
+                    self.get_last_jumpable_mut_between(1, if_block_distance as usize)
+                {
                     skippable.set_c_to_boolean(true);
                 } else {
                     if_block_distance += 1;
@@ -1173,7 +1181,6 @@ impl<'src> Compiler<'src> {
                     );
                 }
             }
-            1 => {}
             2.. => {
                 if_block_distance += 1;
 
diff --git a/dust-lang/tests/control_flow.rs b/dust-lang/tests/control_flow.rs
index 6c83b86..733bb05 100644
--- a/dust-lang/tests/control_flow.rs
+++ b/dust-lang/tests/control_flow.rs
@@ -333,18 +333,20 @@ fn if_else_false() {
                         .set_c_is_constant(),
                     Span(5, 7)
                 ),
-                (Instruction::jump(1, true), Span(10, 11)),
+                (Instruction::jump(2, true), Span(10, 11)),
                 (
                     Instruction::call_native(0, NativeFunction::Panic, 0),
                     Span(12, 19)
                 ),
-                (Instruction::load_constant(1, 2, true), Span(29, 31)),
-                (Instruction::r#move(1, 0), Span(33, 33)),
-                (Instruction::r#return(true), Span(33, 33)),
+                (Instruction::load_constant(0, 2, true), Span(21, 22)),
+                (Instruction::load_constant(1, 3, false), Span(32, 34)),
+                (Instruction::r#move(1, 0), Span(36, 36)),
+                (Instruction::r#return(true), Span(36, 36)),
             ],
             vec![
                 ConcreteValue::Integer(1),
                 ConcreteValue::Integer(2),
+                ConcreteValue::Integer(0),
                 ConcreteValue::Integer(42)
             ],
             vec![]
@@ -374,16 +376,22 @@ fn if_else_true() {
                         .set_c_is_constant(),
                     Span(5, 7)
                 ),
-                (Instruction::jump(1, true), Span(10, 11)),
-                (Instruction::load_constant(0, 1, true), Span(12, 14)),
+                (Instruction::jump(2, true), Span(10, 11)),
+                (Instruction::load_constant(0, 1, false), Span(12, 14)),
+                (Instruction::jump(2, true), Span(36, 36)),
                 (
-                    Instruction::call_native(1, NativeFunction::Panic, 0),
+                    Instruction::call_native(0, NativeFunction::Panic, 0),
                     Span(24, 31)
                 ),
-                (Instruction::r#move(1, 0), Span(33, 33)),
-                (Instruction::r#return(true), Span(33, 33))
+                (Instruction::load_constant(1, 2, false), Span(33, 34)),
+                (Instruction::r#move(1, 0), Span(36, 36)),
+                (Instruction::r#return(true), Span(36, 36))
+            ],
+            vec![
+                ConcreteValue::Integer(1),
+                ConcreteValue::Integer(42),
+                ConcreteValue::Integer(0)
             ],
-            vec![ConcreteValue::Integer(1), ConcreteValue::Integer(42)],
             vec![]
         )),
     );