Begin fixing control flow

dust-lang/src/optimizer.rs

@@ -39,7 +39,7 @@ impl<'chunk> Optimizer<'chunk> {
     }
 
     fn optimize_comparison(&mut self) {
-        log::trace!("Optimizing comparison");
+        log::debug!("Optimizing comparison");
 
         let first_loader_register = {
             let first_loader = &mut self.instructions[2].0;

dust-lang/src/parser.rs

@@ -252,19 +252,6 @@ impl<'src> Parser<'src> {
         Some(operations)
     }
 
-    fn get_last_jump_mut(&mut self) -> Option<&mut Instruction> {
-        self.chunk
-            .instructions_mut()
-            .iter_mut()
-            .find_map(|(instruction, _)| {
-                if let Operation::Jump = instruction.operation() {
-                    Some(instruction)
-                } else {
-                    None
-                }
-            })
-    }
-
     fn get_last_jumpable_mut(&mut self) -> Option<&mut Instruction> {
         self.chunk
             .instructions_mut()
@@ -901,12 +888,18 @@ impl<'src> Parser<'src> {
         }
     }
 
-    fn parse_block(&mut self, allowed: Allowed) -> Result<(), ParseError> {
+    fn parse_block(&mut self, _: Allowed) -> Result<(), ParseError> {
         self.advance()?;
         self.chunk.begin_scope();
 
         while !self.allow(Token::RightCurlyBrace)? && !self.is_eof() {
-            self.parse(Precedence::None, allowed)?;
+            self.parse(
+                Precedence::None,
+                Allowed {
+                    assignment: true,
+                    explicit_return: true,
+                },
+            )?;
         }
 
         self.chunk.end_scope();
@@ -962,28 +955,18 @@ impl<'src> Parser<'src> {
                 Operation::LoadBoolean,
                 Operation::Jump,
                 Operation::Equal | Operation::Less | Operation::LessEqual
-            ],)
+            ])
         ) {
             self.chunk.instructions_mut().pop();
             self.chunk.instructions_mut().pop();
+            self.chunk.instructions_mut().pop();
         }
 
-        let block_allowed = Allowed {
+        let if_block_start = self.chunk.len();
-            assignment: allowed.assignment,
+        let if_block_start_position = self.current_position;
-            explicit_return: allowed.explicit_return,
-        };
 
         if let Token::LeftCurlyBrace = self.current_token {
-            let block_start = self.chunk.len();
+            self.parse_block(allowed)?;
-
-            self.parse_block(block_allowed)?;
-
-            let block_end = self.chunk.len();
-            let jump_distance = (block_end - block_start) as u8;
-
-            if let Some(jump) = self.get_last_jump_mut() {
-                jump.set_b(jump_distance);
-            }
         } else {
             return Err(ParseError::ExpectedToken {
                 expected: TokenKind::LeftCurlyBrace,
@@ -992,6 +975,9 @@ impl<'src> Parser<'src> {
             });
         }
 
+        let if_block_end = self.chunk.len();
+        let mut if_block_distance = (if_block_end - if_block_start) as u8;
+
         let if_block_is_expression = self
             .chunk
             .instructions()
@@ -1006,11 +992,66 @@ impl<'src> Parser<'src> {
             })
             .unwrap_or(false);
 
-        if let Token::Else = self.current_token {
-            let else_start = self.chunk.len();
         let if_last_register = self.next_register().saturating_sub(1);
 
-            self.parse_else(allowed, block_allowed)?;
+        if let Token::Else = self.current_token {
+            self.advance()?;
+
+            if let Token::LeftCurlyBrace = self.current_token {
+                self.parse_block(allowed)?;
+            } else {
+                return Err(ParseError::ExpectedTokenMultiple {
+                    expected: &[TokenKind::If, TokenKind::LeftCurlyBrace],
+                    found: self.current_token.to_owned(),
+                    position: self.current_position,
+                });
+            }
+        } else {
+            self.previous_is_expression = false;
+        }
+
+        self.previous_is_expression = if_block_is_expression && self.previous_is_expression;
+
+        let else_block_end = self.chunk.len();
+        let else_block_distance = (else_block_end - if_block_end) as u8;
+
+        match else_block_distance {
+            0 => {}
+            1 => {
+                if let Some(skippable) = self.get_last_jumpable_mut() {
+                    skippable.set_c_to_boolean(true);
+                } else {
+                    if_block_distance += 1;
+
+                    self.chunk.instructions_mut().insert(
+                        if_block_end,
+                        (
+                            Instruction::jump(else_block_distance, true),
+                            self.current_position,
+                        ),
+                    );
+                }
+            }
+            2.. => {
+                if_block_distance += 1;
+
+                self.chunk.instructions_mut().insert(
+                    if_block_end,
+                    (
+                        Instruction::jump(else_block_distance, true),
+                        self.current_position,
+                    ),
+                );
+            }
+        }
+
+        self.chunk.instructions_mut().insert(
+            if_block_start,
+            (
+                Instruction::jump(if_block_distance, true),
+                if_block_start_position,
+            ),
+        );
 
         if self.chunk.len() >= 4 {
             let possible_comparison_statement = {
@@ -1023,8 +1064,6 @@ impl<'src> Parser<'src> {
         }
 
         let else_last_register = self.next_register().saturating_sub(1);
-            let else_end = self.chunk.len();
-            let jump_distance = (else_end - else_start) as u8;
 
         if if_last_register < else_last_register {
             self.emit_instruction(
@@ -1033,67 +1072,6 @@ impl<'src> Parser<'src> {
             );
         }
 
-            self.previous_is_expression = if_block_is_expression && self.previous_is_expression;
-
-            if jump_distance == 1 {
-                if let Some(skippable) = self.get_last_jumpable_mut() {
-                    skippable.set_c_to_boolean(true);
-                } else {
-                    self.chunk.instructions_mut().insert(
-                        else_start,
-                        (
-                            Instruction::jump(jump_distance, true),
-                            self.current_position,
-                        ),
-                    );
-                }
-            } else {
-                self.chunk.instructions_mut().insert(
-                    else_start,
-                    (
-                        Instruction::jump(jump_distance, true),
-                        self.current_position,
-                    ),
-                );
-            }
-        } else {
-            self.previous_is_expression = false;
-        }
-
-        Ok(())
-    }
-
-    fn parse_else(&mut self, allowed: Allowed, block_allowed: Allowed) -> Result<(), ParseError> {
-        self.advance()?;
-
-        let if_block_end = self.chunk.len();
-
-        if let Token::If = self.current_token {
-            self.parse_if(allowed)?;
-        } else if let Token::LeftCurlyBrace = self.current_token {
-            self.parse_block(block_allowed)?;
-
-            let else_end = self.chunk.len();
-
-            if else_end - if_block_end > 1 {
-                let jump_distance = (else_end - if_block_end) as u8;
-
-                self.chunk.instructions_mut().insert(
-                    if_block_end,
-                    (
-                        Instruction::jump(jump_distance, true),
-                        self.current_position,
-                    ),
-                );
-            }
-        } else {
-            return Err(ParseError::ExpectedTokenMultiple {
-                expected: &[TokenKind::If, TokenKind::LeftCurlyBrace],
-                found: self.current_token.to_owned(),
-                position: self.current_position,
-            });
-        }
-
         Ok(())
     }
 
@@ -1229,7 +1207,7 @@ impl<'src> Parser<'src> {
         )
     }
 
-    fn parse_return(&mut self, allowed: Allowed) -> Result<(), ParseError> {
+    fn parse_return_statement(&mut self, allowed: Allowed) -> Result<(), ParseError> {
         if !allowed.explicit_return {
             return Err(ParseError::UnexpectedReturn {
                 position: self.current_position,
@@ -1260,9 +1238,7 @@ impl<'src> Parser<'src> {
     }
 
     fn parse_implicit_return(&mut self) -> Result<(), ParseError> {
-        let has_semicolon = self.allow(Token::Semicolon)?;
+        if self.allow(Token::Semicolon)? {
-
-        if has_semicolon {
             self.emit_instruction(Instruction::r#return(false), self.current_position);
         } else {
             self.emit_instruction(
@@ -1816,7 +1792,7 @@ impl From<&Token<'_>> for ParseRule<'_> {
                 precedence: Precedence::Assignment,
             },
             Token::Return => ParseRule {
-                prefix: Some(Parser::parse_return),
+                prefix: Some(Parser::parse_return_statement),
                 infix: None,
                 precedence: Precedence::None,
             },

dust-lang/tests/control_flow.rs

@@ -15,7 +15,7 @@ fn equality_assignment_long() {
                         .set_c_is_constant(),
                     Span(13, 15)
                 ),
-                (Instruction::jump(1, true), Span(13, 15)),
+                (Instruction::jump(1, true), Span(18, 19)),
                 (Instruction::load_boolean(0, true, true), Span(20, 24)),
                 (Instruction::load_boolean(0, false, false), Span(34, 39)),
                 (Instruction::define_local(0, 0, false), Span(4, 5)),
@@ -61,7 +61,65 @@ fn equality_assignment_short() {
 }
 
 #[test]
-fn if_else_assigment() {
+fn if_else_assigment_false() {
+    let source = r#"
+        let a = if 4 == 3 {
+            1; 2; 3; 4;
+            panic()
+        } else {
+            1; 2; 3; 4;
+            42
+        };
+        a"#;
+
+    assert_eq!(
+        parse(source),
+        Ok(Chunk::with_data(
+            None,
+            vec![
+                (
+                    *Instruction::equal(true, 0, 1)
+                        .set_b_is_constant()
+                        .set_c_is_constant(),
+                    Span(22, 24)
+                ),
+                (Instruction::jump(6, true), Span(27, 28)),
+                (Instruction::load_constant(0, 2, false), Span(41, 42)),
+                (Instruction::load_constant(1, 3, false), Span(44, 45)),
+                (Instruction::load_constant(2, 1, false), Span(47, 48)),
+                (Instruction::load_constant(3, 0, false), Span(50, 51)),
+                (
+                    Instruction::call_native(4, NativeFunction::Panic, 0),
+                    Span(65, 72)
+                ),
+                (Instruction::jump(5, true), Span(138, 139)),
+                (Instruction::load_constant(5, 2, false), Span(102, 103)),
+                (Instruction::load_constant(6, 3, false), Span(105, 106)),
+                (Instruction::load_constant(7, 1, false), Span(108, 109)),
+                (Instruction::load_constant(8, 0, false), Span(111, 112)),
+                (Instruction::load_constant(9, 4, false), Span(126, 128)),
+                (Instruction::r#move(9, 4), Span(138, 139)),
+                (Instruction::define_local(9, 0, false), Span(13, 14)),
+                (Instruction::get_local(10, 0), Span(148, 149)),
+                (Instruction::r#return(true), Span(149, 149)),
+            ],
+            vec![
+                Value::integer(4),
+                Value::integer(3),
+                Value::integer(1),
+                Value::integer(2),
+                Value::integer(42),
+                Value::string("a")
+            ],
+            vec![Local::new(5, None, false, Scope::default(), 0)]
+        )),
+    );
+
+    assert_eq!(run(source), Ok(Some(Value::integer(42))));
+}
+
+#[test]
+fn if_else_assigment_true() {
     let source = r#"
         let a = if 4 == 4 {
             1; 2; 3; 4;
@@ -83,13 +141,12 @@ fn if_else_assigment() {
                         .set_c_is_constant(),
                     Span(22, 24)
                 ),
-                (Instruction::jump(5, true), Span(22, 24)),
+                (Instruction::jump(6, true), Span(27, 28)),
                 (Instruction::load_constant(0, 1, false), Span(41, 42)),
                 (Instruction::load_constant(1, 2, false), Span(44, 45)),
                 (Instruction::load_constant(2, 3, false), Span(47, 48)),
                 (Instruction::load_constant(3, 0, false), Span(50, 51)),
                 (Instruction::load_constant(4, 4, false), Span(65, 67)),
-                (Instruction::jump(6, true), Span(138, 139)),
                 (Instruction::jump(5, true), Span(138, 139)),
                 (Instruction::load_constant(5, 1, false), Span(97, 98)),
                 (Instruction::load_constant(6, 2, false), Span(100, 101)),
@@ -139,12 +196,11 @@ fn if_else_complex() {
                         .set_c_is_constant(),
                     Span(14, 16)
                 ),
-                (Instruction::jump(5, true), Span(14, 16)),
+                (Instruction::jump(5, true), Span(19, 20)),
                 (Instruction::load_constant(0, 0, false), Span(33, 34)),
                 (Instruction::load_constant(1, 1, false), Span(36, 37)),
                 (Instruction::load_constant(2, 2, false), Span(39, 40)),
                 (Instruction::load_constant(3, 3, false), Span(42, 43)),
-                (Instruction::jump(5, true), Span(95, 95)),
                 (Instruction::jump(4, true), Span(95, 95)),
                 (Instruction::load_constant(4, 0, false), Span(74, 75)),
                 (Instruction::load_constant(5, 1, false), Span(77, 78)),
@@ -251,7 +307,7 @@ fn if_else_false() {
                         .set_c_is_constant(),
                     Span(5, 7)
                 ),
-                (Instruction::jump(1, true), Span(5, 7)),
+                (Instruction::jump(1, true), Span(10, 11)),
                 (
                     Instruction::call_native(0, NativeFunction::Panic, 0),
                     Span(12, 19)
@@ -283,7 +339,7 @@ fn if_else_true() {
                         .set_c_is_constant(),
                     Span(5, 7)
                 ),
-                (Instruction::jump(1, true), Span(5, 7)),
+                (Instruction::jump(1, true), Span(10, 11)),
                 (Instruction::load_constant(0, 1, true), Span(12, 14)),
                 (
                     Instruction::call_native(1, NativeFunction::Panic, 0),
@@ -315,7 +371,7 @@ fn if_false() {
                         .set_c_is_constant(),
                     Span(5, 7)
                 ),
-                (Instruction::jump(1, true), Span(5, 7)),
+                (Instruction::jump(1, true), Span(10, 11)),
                 (Instruction::load_constant(0, 1, false), Span(12, 13)),
                 (Instruction::r#return(false), Span(15, 15))
             ],
@@ -342,7 +398,7 @@ fn if_true() {
                         .set_c_is_constant(),
                     Span(5, 7)
                 ),
-                (Instruction::jump(1, true), Span(5, 7)),
+                (Instruction::jump(1, true), Span(10, 11)),
                 (Instruction::load_constant(0, 1, false), Span(12, 13)),
                 (Instruction::r#return(false), Span(15, 15))
             ],

examples/fibonacci.ds

@@ -1,10 +1,8 @@
 fn fib (n: int) -> int {
-    if n <= 0 { 0 }
+    if n <= 0 { return 0 }
-    if n == 1 {
+    if n == 1 { return 1 }
-        1
+
-    } else {
     fib(n - 1) + fib(n - 2)
 }
-}
 
 fib(10)