Fix control flow

dust-lang/src/parser/mod.rs

@@ -542,6 +542,7 @@ impl<'src> Parser<'src> {
             Instruction::load_boolean(self.current_register, false, false),
             operator_position,
         );
+        self.increment_register()?;
 
         Ok(())
     }
@@ -730,10 +731,32 @@ impl<'src> Parser<'src> {
         self.advance()?;
         self.parse_expression()?;
 
+        if matches!(
+            self.chunk.get_last_n_operations(),
+            [
+                Some(Operation::LoadBoolean),
+                Some(Operation::LoadBoolean),
+                Some(Operation::Jump)
+            ]
+        ) {
+            self.chunk.pop_instruction(self.current_position)?;
+            self.chunk.pop_instruction(self.current_position)?;
+            self.chunk.pop_instruction(self.current_position)?;
+            self.decrement_register()?;
+        }
+
+        let jump_start = self.chunk.len();
+
         if let Token::LeftCurlyBrace = self.current_token {
             self.parse_block(allow_assignment, allow_return)?;
         }
 
+        let jump_end = self.chunk.len();
+        let jump_instruction = Instruction::jump((jump_end - jump_start) as u8, true);
+
+        self.chunk
+            .insert_instruction(jump_start, jump_instruction, self.current_position);
+
         if self.allow(TokenKind::Else)? {
             if let Token::If = self.current_token {
                 self.parse_if(allow_assignment, allow_return)?;

dust-lang/src/parser/tests.rs

@@ -366,10 +366,10 @@ fn block_scope() {
                 (Instruction::define_local(1, 1, false), Span(46, 47)),
                 (Instruction::load_constant(2, 2), Span(92, 93)),
                 (Instruction::define_local(2, 2, false), Span(88, 89)),
-                (Instruction::end(false), Span(121, 124)),
+                (Instruction::end(false), Span(107, 108)),
                 (Instruction::load_constant(3, 3), Span(129, 130)),
                 (Instruction::define_local(3, 3, false), Span(125, 126)),
-                (Instruction::end(false), Span(150, 159)),
+                (Instruction::end(false), Span(140, 141)),
                 (Instruction::load_constant(4, 4), Span(158, 159)),
                 (Instruction::define_local(4, 4, false), Span(154, 155)),
                 (Instruction::end(false), Span(165, 165)),

dust-lang/src/vm.rs

@@ -1,3 +1,5 @@
+use log::debug;
+
 use crate::{
     parse, AnnotatedError, Chunk, ChunkError, DustError, Identifier, Instruction, Local, Operation,
     Span, Value, ValueError,
@@ -60,7 +62,7 @@ impl Vm {
         while let Ok((instruction, position)) = self.read(Span(0, 0)).copied() {
             log::trace!(
                 "Running IP {} {} at {position}",
-                self.ip,
+                self.ip - 1,
                 instruction.operation()
             );
 
@@ -220,42 +222,96 @@ impl Vm {
                     }
                 }
                 Operation::Equal => {
+                    let (jump, _) = *self.chunk.get_instruction(self.ip, position)?;
+
+                    debug_assert_eq!(jump.operation(), Operation::Jump);
+
                     let (left, right) = get_arguments(self, instruction, position)?;
-                    let equal = left
+                    let boolean = left
                         .equal(right)
-                        .map_err(|error| VmError::Value { error, position })?;
+                        .map_err(|error| VmError::Value { error, position })?
+                        .as_boolean()
+                        .ok_or_else(|| VmError::ExpectedBoolean {
+                            found: left.clone(),
+                            position,
+                        })?;
                     let compare_to = instruction.destination_as_boolean();
 
-                    if let Some(boolean) = equal.as_boolean() {
+                    if boolean == compare_to {
-                        if boolean && compare_to {
                         self.ip += 1;
-                        }
+                    } else {
+                        let jump_distance = jump.destination();
+                        let is_positive = jump.first_argument_as_boolean();
+                        let new_ip = if is_positive {
+                            self.ip + jump_distance as usize
+                        } else {
+                            self.ip - jump_distance as usize
+                        };
+
+                        self.ip = new_ip;
                     }
                 }
                 Operation::Less => {
+                    let (jump, _) = *self.chunk.get_instruction(self.ip, position)?;
+
+                    assert_eq!(jump.operation(), Operation::Jump);
+
                     let (left, right) = get_arguments(self, instruction, position)?;
                     let less = left
                         .less_than(right)
                         .map_err(|error| VmError::Value { error, position })?;
+                    let boolean = left
+                        .equal(right)
+                        .map_err(|error| VmError::Value { error, position })?
+                        .as_boolean()
+                        .ok_or_else(|| VmError::ExpectedBoolean {
+                            found: less,
+                            position,
+                        })?;
                     let compare_to = instruction.destination_as_boolean();
 
-                    if let Some(boolean) = less.as_boolean() {
                     if boolean == compare_to {
                         self.ip += 1;
-                        }
+                    } else {
+                        let jump_distance = jump.destination();
+                        let is_positive = jump.first_argument_as_boolean();
+                        let new_ip = if is_positive {
+                            self.ip + jump_distance as usize
+                        } else {
+                            self.ip - jump_distance as usize
+                        };
+
+                        self.ip = new_ip;
                     }
                 }
                 Operation::LessEqual => {
+                    let (jump, _) = *self.read(position)?;
                     let (left, right) = get_arguments(self, instruction, position)?;
-                    let less_equal = left
+                    let less_or_equal = left
                         .less_than_or_equal(right)
                         .map_err(|error| VmError::Value { error, position })?;
+                    let boolean = left
+                        .equal(right)
+                        .map_err(|error| VmError::Value { error, position })?
+                        .as_boolean()
+                        .ok_or_else(|| VmError::ExpectedBoolean {
+                            found: less_or_equal,
+                            position,
+                        })?;
                     let compare_to = instruction.destination_as_boolean();
 
-                    if let Some(boolean) = less_equal.as_boolean() {
                     if boolean == compare_to {
                         self.ip += 1;
-                        }
+                    } else {
+                        let jump_distance = jump.destination();
+                        let is_positive = jump.first_argument_as_boolean();
+                        let new_ip = if is_positive {
+                            self.ip + jump_distance as usize
+                        } else {
+                            self.ip - jump_distance as usize
+                        };
+
+                        self.ip = new_ip;
                     }
                 }
                 Operation::Negate => {
@@ -307,11 +363,9 @@ impl Vm {
                 Operation::End => {
                     let returns_value = instruction.destination_as_boolean();
 
-                    return if returns_value {
+                    if returns_value {
-                        Ok(Some(self.pop(position)?))
+                        return Ok(Some(self.pop(position)?));
-                    } else {
+                    }
-                        Ok(None)
-                    };
                 }
             }
         }