Add if/elseif/else

examples/test.lua

@@ -36,4 +36,24 @@ print(min(11, 9))
 print(10 - 15)
 print("hello there!")
 print(true or 0)
-print(123, f(10, 11, 12))
+
+global value, table = f(10, 11, 12)
+
+print("hello")
+for i=1,#table do
+    print(table[i])
+    if i > 2 then
+        goto test
+    end
+end
+::test::
+
+local test = table[1]
+if test == 10 then
+    print("first")
+elseif test == 11 then
+    print("second")
+else
+    print("third")
+end
+print("after")

src/ast.rs

@@ -234,7 +234,12 @@ impl Parse for Block {
 
         while !matches!(
             stream.peek(),
-            Some(Token::Keyword(Keyword::End) | Token::Eof)
+            Some(
+                Token::Keyword(Keyword::End)
+                    | Token::Keyword(Keyword::ElseIf)
+                    | Token::Keyword(Keyword::Else)
+                    | Token::Eof
+            )
         ) {
             statements.push(stream.parse()?);
         }
@@ -254,8 +259,21 @@ pub enum Statement {
         ExpressionList,
     ),
     Return(ExpressionList),
-    If(Node<Expression>, Block),
+    If {
+        if_part: (Node<Expression>, Block),
+        else_block: Block,
+    },
     Expression(Node<Expression>),
+    NumericalFor(
+        Node<String>,
+        Node<Expression>,
+        Node<Expression>,
+        Node<Expression>,
+        Block,
+    ),
+    Break,
+    Label(Node<String>),
+    GoTo(Node<String>),
 }
 
 impl Parse for Statement {
@@ -286,11 +304,48 @@ impl Parse for Statement {
             Ok(Statement::Return(stream.parse()?))
         } else if peeked == Some(Token::Keyword(Keyword::If)) {
             stream.next(); // Consume if
-            let cond = stream.parse()?;
+            let if_cond = stream.parse()?;
             stream.expect(Token::Keyword(Keyword::Then))?;
-            let then = stream.parse()?;
+            let if_block = stream.parse()?;
+
+            let mut elif_statements = Vec::new();
+            while let Some(Token::Keyword(Keyword::ElseIf)) = stream.peek() {
+                stream.next();
+                let elif_cond: Node<Expression> = stream.parse()?;
+                stream.expect(Token::Keyword(Keyword::Then))?;
+                let elif_block: Block = stream.parse()?;
+                elif_statements.push((elif_cond, elif_block))
+            }
+
+            let mut else_block = if let Some(Token::Keyword(Keyword::Else)) = stream.peek() {
+                stream.next();
+                stream.parse()?
+            } else {
+                Block {
+                    statements: Vec::new(),
+                    _meta: Metadata::empty(),
+                }
+            };
+
+            for elif_statement in elif_statements.into_iter().rev() {
+                let elif_meta = elif_statement.0.meta.clone() + elif_statement.1._meta.clone();
+                else_block = Block {
+                    statements: vec![Node {
+                        kind: Statement::If {
+                            if_part: elif_statement,
+                            else_block: else_block.clone(),
+                        },
+                        meta: elif_meta.clone(),
+                    }],
+                    _meta: elif_meta + else_block._meta,
+                }
+            }
+
             stream.expect(Token::Keyword(Keyword::End))?;
-            Ok(Self::If(cond, then))
+            Ok(Self::If {
+                if_part: (if_cond, if_block),
+                else_block,
+            })
         } else if let Some(Token::Keyword(Keyword::Local | Keyword::Global)) = peeked {
             let access_modifier = match stream.next() {
                 Some(Token::Keyword(Keyword::Local)) => AccessModifier::Local,
@@ -316,8 +371,52 @@ impl Parse for Statement {
                 vec![(access.0, access.1.0)],
                 ExpressionList(vec![expression]),
             ))
+        } else if stream.peek() == Some(Token::Keyword(Keyword::For)) {
+            stream.next();
+            let counter_name = stream.parse()?;
+            stream.expect_symbol('=')?;
+            let init = stream.parse()?;
+            stream.expect_symbol(',')?;
+            let end = stream.parse()?;
+
+            let step = if let Some(Token::Symbol(',')) = stream.peek() {
+                stream.next();
+                stream.parse()?
+            } else {
+                Node {
+                    kind: Expression::Literal(Literal::Integer(LuaInteger(1))),
+                    meta: Metadata::empty(),
+                }
+            };
+
+            stream.expect(Token::Keyword(Keyword::Do))?;
+            let block = stream.parse()?;
+            stream.expect(Token::Keyword(Keyword::End))?;
+
+            Ok(Statement::NumericalFor(
+                counter_name,
+                init,
+                end,
+                step,
+                block,
+            ))
         } else if let Ok(expr) = stream.parse() {
             Ok(Self::Expression(expr))
+        } else if let Some(Token::Keyword(Keyword::Break)) = stream.peek() {
+            stream.next();
+            Ok(Self::Break)
+        } else if let (Some(Token::Symbol(':')), Some(Token::Symbol(':'))) =
+            (stream.peek(), stream.peek2())
+        {
+            stream.next();
+            stream.next();
+            let name = stream.parse()?;
+            stream.expect_symbol(':')?;
+            stream.expect_symbol(':')?;
+            Ok(Self::Label(name))
+        } else if let Some(Token::Keyword(Keyword::GoTo)) = stream.peek() {
+            stream.next();
+            Ok(Self::GoTo(stream.parse()?))
         } else {
             Err(stream.expecting_err("statement"))
         }

src/compile.rs

@@ -47,6 +47,22 @@ impl LocalCounter {
     }
 
     pub fn consecutive(&mut self, amount: usize) -> Vec<u16> {
+        'outer: for free_num in self.1.clone() {
+            let mut potentials = vec![free_num];
+            let mut curr = free_num;
+            for _ in 0..amount {
+                if let Some(next) = self.1.iter().find(|v| **v == curr + 1) {
+                    potentials.push(*next);
+                    curr = *next;
+                } else {
+                    continue 'outer;
+                }
+            }
+
+            self.1
+                .retain_mut(|v| potentials.iter().find(|p| v != *p).is_none());
+            return potentials;
+        }
         let mut returned = Vec::new();
         for _ in 0..amount {
             returned.push(self.new());
@@ -61,6 +77,46 @@ impl LocalCounter {
     }
 }
 
+pub(crate) enum PreInstr {
+    Instr(Instruction),
+    Break,
+    Label(String),
+    GoTo(String),
+}
+
+pub(crate) fn process_pre_instrs(pre_instructions: Vec<PreInstr>) -> Vec<Instruction> {
+    let mut instructions = Vec::new();
+
+    let mut label_positions = HashMap::new();
+
+    let mut counter: u32 = 0;
+    for pre_instr in &pre_instructions {
+        match pre_instr {
+            PreInstr::Label(name) => {
+                label_positions.insert(name.clone(), counter);
+                continue;
+            }
+            _ => {}
+        }
+        counter += 1;
+    }
+
+    for pre_instr in pre_instructions {
+        match pre_instr {
+            PreInstr::Instr(instruction) => instructions.push(instruction),
+            PreInstr::Break => panic!(),
+            PreInstr::GoTo(name) => {
+                instructions.push(Instruction::GoTo(*label_positions.get(&name).unwrap()))
+            }
+            PreInstr::Label(_) => {
+                continue;
+            }
+        }
+    }
+
+    instructions
+}
+
 impl Block {
     pub(crate) fn find_constants(
         &self,
@@ -78,7 +134,7 @@ impl Block {
         constants
     }
 
-    pub(crate) fn compile(&self, state: &mut State, scope: &mut Scope) -> Vec<Instruction> {
+    pub(crate) fn compile(&self, state: &mut State, scope: &mut Scope) -> Vec<PreInstr> {
         let mut instructions = Vec::new();
 
         let mut inner_scope = scope.clone();
@@ -131,17 +187,33 @@ impl Statement {
                 }
                 constants
             }
-            Statement::If(cond, then) => {
+            Statement::If {
+                if_part,
+                else_block,
+            } => {
                 let mut constants = HashSet::new();
-                constants.extend(cond.kind.find_constants(scope));
-                constants.extend(then.find_constants(scope, Vec::new()));
+                constants.extend(if_part.0.kind.find_constants(scope));
+                constants.extend(if_part.1.find_constants(scope, Vec::new()));
+
+                constants.extend(else_block.find_constants(scope, Vec::new()));
                 constants
             }
             Statement::Expression(expr) => expr.kind.find_constants(scope),
+            Statement::NumericalFor(_, init, end, step, block) => {
+                let mut constants = HashSet::new();
+                constants.extend(init.kind.find_constants(scope));
+                constants.extend(end.kind.find_constants(scope));
+                constants.extend(step.kind.find_constants(scope));
+                constants.extend(block.find_constants(scope, Vec::new()));
+                constants
+            }
+            Statement::Break => HashSet::new(),
+            Statement::Label(_) => HashSet::new(),
+            Statement::GoTo(_) => HashSet::new(),
         }
     }
 
-    fn compile(&self, state: &mut State, scope: &mut Scope) -> Vec<Instruction> {
+    fn compile(&self, state: &mut State, scope: &mut Scope) -> Vec<PreInstr> {
         let mut instructions = Vec::new();
 
         match self {
@@ -149,7 +221,7 @@ impl Statement {
                 let new_registers = if *access_modifier == Some(AccessModifier::Local) {
                     let min_reg = scope.register_counter.0 + 1;
                     let max_reg = scope.register_counter.0 + names.len() as u16;
-                    instructions.push(Instruction::LoadNil(min_reg, max_reg));
+                    instructions.push(PreInstr::Instr(Instruction::LoadNil(min_reg, max_reg)));
                     scope.register_counter.0 += names.len() as u16 + 1;
 
                     let mut new_registers = Vec::new();
@@ -181,14 +253,14 @@ impl Statement {
                         let mut vararg_applied = false;
                         for (i, (name, indexes)) in names.iter().enumerate() {
                             if let Some(expr_reg) = expr_regs.get(i) {
-                                instructions.push(Instruction::Move(
+                                instructions.push(PreInstr::Instr(Instruction::Move(
                                     *new_registers.get(i).unwrap(),
                                     *expr_reg,
-                                ));
+                                )));
                             } else if !vararg_applied {
-                                instructions.push(Instruction::MoveRetValues(
+                                instructions.push(PreInstr::Instr(Instruction::MoveRetValues(
                                     *new_registers.get(i).unwrap(),
-                                ));
+                                )));
                                 vararg_applied = true;
                             }
                             if indexes.len() > 0 {
@@ -205,10 +277,10 @@ impl Statement {
                                 todo!()
                             }
                             let global = state.get_constant(&Constant::String(name.kind.clone()));
-                            instructions.push(Instruction::SetGlobal(
+                            instructions.push(PreInstr::Instr(Instruction::SetGlobal(
                                 expr_regs.get(i).cloned().unwrap(),
                                 global,
-                            ));
+                            )));
                         }
                     }
                     None => {
@@ -218,13 +290,17 @@ impl Statement {
                                     *reg
                                 } else if let Some(upval_reg) = scope.upvalues.get(&name.kind) {
                                     let local = scope.register_counter.next();
-                                    instructions.push(Instruction::GetUpVal(local, *upval_reg));
+                                    instructions.push(PreInstr::Instr(Instruction::GetUpVal(
+                                        local, *upval_reg,
+                                    )));
                                     local
                                 } else {
                                     let global =
                                         state.get_constant(&Constant::String(name.kind.clone()));
                                     let local = scope.register_counter.next();
-                                    instructions.push(Instruction::GetGlobal(local, global));
+                                    instructions.push(PreInstr::Instr(Instruction::GetGlobal(
+                                        local, global,
+                                    )));
                                     local
                                 };
 
@@ -237,40 +313,40 @@ impl Statement {
                                         let (instr, index_reg) =
                                             index.kind.compile(state, scope, Some(1));
                                         instructions.extend(instr);
-                                        instructions.push(Instruction::GetTable(
+                                        instructions.push(PreInstr::Instr(Instruction::GetTable(
                                             table_reg,
                                             table_reg,
                                             *index_reg.first().unwrap(),
-                                        ));
+                                        )));
                                     }
                                 }
 
                                 let (instr, index_reg) =
                                     indexes.last().unwrap().kind.compile(state, scope, Some(1));
                                 instructions.extend(instr);
-                                instructions.push(Instruction::SetTable(
+                                instructions.push(PreInstr::Instr(Instruction::SetTable(
                                     table_reg,
                                     *index_reg.first().unwrap(),
                                     expr_regs.get(i).cloned().unwrap(),
-                                ));
+                                )));
                             } else {
                                 if let Some(reg) = scope.locals.get(&name.kind) {
-                                    instructions.push(Instruction::Move(
+                                    instructions.push(PreInstr::Instr(Instruction::Move(
                                         *reg,
                                         expr_regs.get(i).cloned().unwrap(),
-                                    ));
+                                    )));
                                 } else if let Some(upval_reg) = scope.upvalues.get(&name.kind) {
-                                    instructions.push(Instruction::SetUpVal(
+                                    instructions.push(PreInstr::Instr(Instruction::SetUpVal(
                                         *upval_reg,
                                         expr_regs.get(i).cloned().unwrap(),
-                                    ));
+                                    )));
                                 } else {
                                     let global =
                                         state.get_constant(&Constant::String(name.kind.clone()));
-                                    instructions.push(Instruction::SetGlobal(
+                                    instructions.push(PreInstr::Instr(Instruction::SetGlobal(
                                         expr_regs.get(i).cloned().unwrap(),
                                         global,
-                                    ));
+                                    )));
                                 }
                             }
                         }
@@ -305,43 +381,97 @@ impl Statement {
                 for (i, ret_register) in ret_registers.iter_mut().enumerate() {
                     let new_reg = first_ret_register + i as u16;
                     if *ret_register != new_reg {
-                        instructions.push(Instruction::Move(new_reg, *ret_register));
+                        instructions
+                            .push(PreInstr::Instr(Instruction::Move(new_reg, *ret_register)));
                     }
                     *ret_register = new_reg;
                 }
 
                 if vararg {
-                    instructions.push(Instruction::MoveRetValues(
+                    instructions.push(PreInstr::Instr(Instruction::MoveRetValues(
                         first_ret_register + ret_registers.len() as u16,
-                    ));
+                    )));
                 }
 
-                dbg!(&first_ret_register);
-
-                instructions.push(Instruction::Return(
+                instructions.push(PreInstr::Instr(Instruction::Return(
                     first_ret_register,
                     if vararg {
                         0
                     } else {
                         *ret_registers.last().unwrap_or(&0)
                     },
-                ));
+                )));
             }
-            Statement::If(expr, block) => {
-                let (instr, regs) = expr.kind.compile(state, scope, Some(1));
+            Statement::If {
+                if_part,
+                else_block,
+            } => {
+                let (instr, regs) = if_part.0.kind.compile(state, scope, Some(1));
                 instructions.extend(instr);
 
                 let local = scope.register_counter.next();
-                instructions.push(Instruction::Test(local, *regs.first().unwrap(), 0));
+                instructions.push(PreInstr::Instr(Instruction::Test(
+                    local,
+                    *regs.first().unwrap(),
+                    0,
+                )));
 
-                let block_instructions = block.compile(state, scope);
-                instructions.push(Instruction::Jmp(block_instructions.len() as i32));
-                instructions.extend(block_instructions);
+                let if_block_instructions = if_part.1.compile(state, scope);
+                let else_block_instructions = else_block.compile(state, scope);
+                instructions.push(PreInstr::Instr(Instruction::Jmp(
+                    if_block_instructions.len() as i32 + 1,
+                )));
+                instructions.extend(if_block_instructions);
+                instructions.push(PreInstr::Instr(Instruction::Jmp(
+                    else_block_instructions.len() as i32,
+                )));
+                instructions.extend(else_block_instructions);
             }
             Statement::Expression(expr) => {
                 let (instr, _) = expr.kind.compile(state, scope, None);
                 instructions.extend(instr);
             }
+            Statement::NumericalFor(counter, init, end, step, block) => {
+                let (instr, init_reg) = init.kind.compile(state, scope, Some(1));
+                instructions.extend(instr);
+                let (instr, end_reg) = end.kind.compile(state, scope, Some(1));
+                instructions.extend(instr);
+                let (instr, step_reg) = step.kind.compile(state, scope, Some(1));
+                instructions.extend(instr);
+
+                let init_reg = init_reg.first().unwrap();
+                let end_reg = end_reg.first().unwrap();
+                let step_reg = step_reg.first().unwrap();
+
+                scope.locals.insert(counter.kind.clone(), *init_reg);
+                scope.locals.insert("_END".to_owned(), *end_reg);
+                scope.locals.insert("_STEP".to_owned(), *step_reg);
+
+                instructions.push(PreInstr::Instr(Instruction::ForTest(
+                    *init_reg, *end_reg, *step_reg,
+                )));
+
+                let mut inner_scope = scope.clone();
+                let instr = block.compile(state, &mut inner_scope);
+                let instr_len = instr.len() as i32;
+
+                instructions.push(PreInstr::Instr(Instruction::Jmp(instr_len + 2)));
+                for (i, pre_instr) in instr.into_iter().enumerate() {
+                    match pre_instr {
+                        PreInstr::Break => instructions
+                            .push(PreInstr::Instr(Instruction::Jmp(instr_len - i as i32 + 1))),
+                        _ => instructions.push(pre_instr),
+                    }
+                }
+
+                instructions.push(PreInstr::Instr(Instruction::Add(
+                    *init_reg, *init_reg, *step_reg,
+                )));
+                instructions.push(PreInstr::Instr(Instruction::Jmp(-(instr_len + 4))));
+            }
+            Statement::Break => instructions.push(PreInstr::Break),
+            Statement::Label(node) => instructions.push(PreInstr::Label(node.kind.clone())),
+            Statement::GoTo(node) => instructions.push(PreInstr::GoTo(node.kind.clone())),
         }
 
         for reg in 0..scope.register_counter.0 {
@@ -442,21 +572,24 @@ impl Expression {
         state: &mut State,
         scope: &mut Scope,
         expected_values: Option<usize>,
-    ) -> (Vec<Instruction>, Vec<u16>) {
+    ) -> (Vec<PreInstr>, Vec<u16>) {
         match self {
             Expression::ValueRef(name) => {
                 if let Some(reg) = scope.locals.get(name) {
                     (Vec::new(), vec![*reg])
                 } else if let Some(upvalue) = scope.upvalues.get(name) {
                     let local = scope.register_counter.next();
-                    (vec![Instruction::GetUpVal(local, *upvalue)], vec![local])
+                    (
+                        vec![PreInstr::Instr(Instruction::GetUpVal(local, *upvalue))],
+                        vec![local],
+                    )
                 } else {
                     let mut instructions = Vec::new();
                     let reg = scope.register_counter.next();
-                    instructions.push(Instruction::GetGlobal(
+                    instructions.push(PreInstr::Instr(Instruction::GetGlobal(
                         reg,
                         state.get_constant(&Constant::String(name.clone())),
-                    ));
+                    )));
                     (instructions, vec![reg])
                 }
             }
@@ -469,64 +602,69 @@ impl Expression {
                 let reg = scope.register_counter.next();
                 match binary_operator {
                     BinaryOperator::Equal => {
-                        instructions.push(Instruction::Equal(
+                        instructions.push(PreInstr::Instr(Instruction::Equal(
                             reg,
                             *lhs.get(0).unwrap(),
                             *rhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                     BinaryOperator::LessThan => {
-                        instructions.push(Instruction::LessThan(
+                        instructions.push(PreInstr::Instr(Instruction::LessThan(
                             reg,
                             *lhs.get(0).unwrap(),
                             *rhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                     BinaryOperator::LessThanOrEqual => {
-                        instructions.push(Instruction::LessThanOrEqual(
+                        instructions.push(PreInstr::Instr(Instruction::LessThanOrEqual(
                             reg,
                             *lhs.get(0).unwrap(),
                             *rhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                     BinaryOperator::GreaterThan => {
-                        instructions.push(Instruction::LessThan(
+                        instructions.push(PreInstr::Instr(Instruction::LessThan(
                             reg,
                             *rhs.get(0).unwrap(),
                             *lhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                     BinaryOperator::GreaterThanOrEqual => {
-                        instructions.push(Instruction::LessThanOrEqual(
+                        instructions.push(PreInstr::Instr(Instruction::LessThanOrEqual(
                             reg,
                             *rhs.get(0).unwrap(),
                             *lhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                     BinaryOperator::Add => {
-                        instructions.push(Instruction::Add(
+                        instructions.push(PreInstr::Instr(Instruction::Add(
                             reg,
                             *lhs.get(0).unwrap(),
                             *rhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                     BinaryOperator::Sub => {
-                        instructions.push(Instruction::Unm(reg, *rhs.get(0).unwrap()));
-                        instructions.push(Instruction::Add(reg, *lhs.get(0).unwrap(), reg));
+                        instructions
+                            .push(PreInstr::Instr(Instruction::Unm(reg, *rhs.get(0).unwrap())));
+                        instructions.push(PreInstr::Instr(Instruction::Add(
+                            reg,
+                            *lhs.get(0).unwrap(),
+                            reg,
+                        )));
                     }
                     BinaryOperator::And => {
-                        instructions.push(Instruction::And(
+                        instructions.push(PreInstr::Instr(Instruction::And(
                             reg,
                             *lhs.get(0).unwrap(),
                             *rhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                     BinaryOperator::Or => {
-                        instructions.push(Instruction::Or(
+                        instructions.push(PreInstr::Instr(Instruction::Or(
                             reg,
                             *lhs.get(0).unwrap(),
                             *rhs.get(0).unwrap(),
-                        ));
+                        )));
                     }
                 };
                 (instructions, vec![reg])
@@ -537,8 +675,12 @@ impl Expression {
                 instructions.extend(instr);
                 for reg in &registers {
                     match op {
-                        UnaryOperator::Negation => instructions.push(Instruction::Unm(*reg, *reg)),
-                        UnaryOperator::Length => instructions.push(Instruction::Len(*reg, *reg)),
+                        UnaryOperator::Negation => {
+                            instructions.push(PreInstr::Instr(Instruction::Unm(*reg, *reg)))
+                        }
+                        UnaryOperator::Length => {
+                            instructions.push(PreInstr::Instr(Instruction::Len(*reg, *reg)))
+                        }
                     }
                 }
                 (instructions, registers)
@@ -568,7 +710,7 @@ impl Expression {
 
                 let instructions = block.compile(state, &mut inner_scope);
                 state.prototypes.push(Prototype {
-                    instructions,
+                    instructions: process_pre_instrs(instructions),
                     parameters: if inner_scope.is_vararg {
                         params.len() - 1
                     } else {
@@ -577,9 +719,14 @@ impl Expression {
                 });
 
                 let mut instructions = Vec::new();
-                instructions.push(Instruction::Close(scope.register_counter.0));
+                instructions.push(PreInstr::Instr(Instruction::Close(
+                    scope.register_counter.0,
+                )));
                 let local = scope.register_counter.next();
-                instructions.push(Instruction::Closure(local, state.prototypes.len() as u32));
+                instructions.push(PreInstr::Instr(Instruction::Closure(
+                    local,
+                    state.prototypes.len() as u32,
+                )));
 
                 (instructions, vec![local])
             }
@@ -626,16 +773,19 @@ impl Expression {
                 for (i, param_reg) in original_param_regs.iter().enumerate().rev() {
                     let new_reg = param_regs.get(i).unwrap();
                     if param_reg != new_reg {
-                        instructions.push(Instruction::Move(*new_reg, *param_reg));
+                        instructions.push(PreInstr::Instr(Instruction::Move(*new_reg, *param_reg)));
                     }
                 }
                 if function_reg != *old_function_reg {
-                    instructions.push(Instruction::Move(function_reg, *old_function_reg));
+                    instructions.push(PreInstr::Instr(Instruction::Move(
+                        function_reg,
+                        *old_function_reg,
+                    )));
                 }
 
                 if vararg {
                     let last_reg = param_regs.last().unwrap_or(&function_reg) + 1;
-                    instructions.push(Instruction::MoveRetValues(last_reg));
+                    instructions.push(PreInstr::Instr(Instruction::MoveRetValues(last_reg)));
                 }
 
                 let last_param_reg = param_regs.last().unwrap_or(&function_reg);
@@ -652,7 +802,7 @@ impl Expression {
                     }
                 }
 
-                instructions.push(Instruction::Call(
+                instructions.push(PreInstr::Instr(Instruction::Call(
                     *&function_reg,
                     if vararg { 0 } else { param_regs.len() as u16 },
                     if return_regs.len() == 0 {
@@ -660,14 +810,14 @@ impl Expression {
                     } else {
                         return_regs.len() as u16 + 1
                     },
-                ));
+                )));
 
                 (instructions, return_regs)
             }
             Expression::Literal(literal) => {
                 let mut instructions = Vec::new();
                 let reg = scope.register_counter.next();
-                instructions.push(Instruction::LoadK(
+                instructions.push(PreInstr::Instr(Instruction::LoadK(
                     reg,
                     state.get_constant(&match literal {
                         Literal::Float(value) => Constant::Float(value.vm_number()),
@@ -676,13 +826,13 @@ impl Expression {
                         Literal::Bool(value) => Constant::Bool(LuaBool(*value)),
                         Literal::Nil => Constant::Nil,
                     }),
-                ));
+                )));
                 (instructions, vec![reg])
             }
             Expression::TableConstructor(entries) => {
                 let mut instructions = Vec::new();
                 let reg = scope.register_counter.next();
-                instructions.push(Instruction::NewTable(reg));
+                instructions.push(PreInstr::Instr(Instruction::NewTable(reg)));
 
                 let mut counter = 1;
                 for (i, (key, value)) in entries.iter().enumerate() {
@@ -691,11 +841,11 @@ impl Expression {
                         instructions.extend(instr);
                         let (instr, value_regs) = value.kind.compile(state, scope, Some(1));
                         instructions.extend(instr);
-                        instructions.push(Instruction::SetTable(
+                        instructions.push(PreInstr::Instr(Instruction::SetTable(
                             reg,
                             *key_regs.first().unwrap(),
                             *value_regs.first().unwrap(),
-                        ));
+                        )));
                     } else {
                         let (instr, value_regs) = value.kind.compile(
                             state,
@@ -710,18 +860,19 @@ impl Expression {
 
                         if value_regs.len() > 0 {
                             let key_reg = scope.register_counter.next();
-                            instructions.push(Instruction::LoadK(
+                            instructions.push(PreInstr::Instr(Instruction::LoadK(
                                 key_reg,
                                 state.get_constant(&Constant::Integer(LuaInteger(counter))),
-                            ));
-                            instructions.push(Instruction::SetTable(
+                            )));
+                            instructions.push(PreInstr::Instr(Instruction::SetTable(
                                 reg,
                                 key_reg,
                                 *value_regs.first().unwrap(),
-                            ));
+                            )));
                             counter += 1;
                         } else {
-                            instructions.push(Instruction::SetList(reg, counter as u32));
+                            instructions
+                                .push(PreInstr::Instr(Instruction::SetList(reg, counter as u32)));
                         }
                     }
                 }
@@ -736,11 +887,11 @@ impl Expression {
                 instructions.extend(instr);
 
                 let local = scope.register_counter.next();
-                instructions.push(Instruction::GetTable(
+                instructions.push(PreInstr::Instr(Instruction::GetTable(
                     local,
                     *expr_regs.first().unwrap(),
                     *index_regs.first().unwrap(),
-                ));
+                )));
 
                 (instructions, vec![local])
             }
@@ -751,10 +902,10 @@ impl Expression {
                 let mut instructions = Vec::new();
                 let new_reg = scope.register_counter.new();
                 if expected_values == None {
-                    instructions.push(Instruction::Vararg(new_reg, 0));
+                    instructions.push(PreInstr::Instr(Instruction::Vararg(new_reg, 0)));
                     (instructions, Vec::new())
                 } else {
-                    instructions.push(Instruction::Vararg(new_reg, 3));
+                    instructions.push(PreInstr::Instr(Instruction::Vararg(new_reg, 3)));
                     (instructions, vec![new_reg])
                 }
             }

src/lib.rs

@@ -31,6 +31,7 @@ use thiserror::Error;
 
 use crate::{
     ast::Block,
+    compile::process_pre_instrs,
     token_stream::{
         TokenStream, TokenStreamError,
         lexer::{Error as TokenizerError, Token, tokenize},
@@ -109,7 +110,7 @@ pub fn compile(
     let chunk = stream.parse::<Block>()?;
     stream.expect(Token::Eof)?;
 
-    // dbg!(&chunk);
+    dbg!(&chunk);
 
     let constants = chunk
         .find_constants(
@@ -139,7 +140,7 @@ pub fn compile(
     // dbg!(&constants);
 
     Ok(CompilationUnit {
-        instructions,
+        instructions: process_pre_instrs(instructions),
         state,
         constants,
     })

src/token_stream/lexer.rs

@@ -15,11 +15,17 @@ pub enum Keyword {
     Global,
     Return,
     If,
+    ElseIf,
+    Else,
     Then,
     True,
     False,
     Nil,
     Not,
+    For,
+    Do,
+    Break,
+    GoTo,
 }
 
 impl Keyword {
@@ -31,11 +37,17 @@ impl Keyword {
             "global" => Keyword::Global,
             "return" => Keyword::Return,
             "if" => Keyword::If,
+            "elseif" => Keyword::ElseIf,
+            "else" => Keyword::Else,
             "then" => Keyword::Then,
             "true" => Keyword::True,
             "false" => Keyword::False,
             "nil" => Keyword::Nil,
             "not" => Keyword::Not,
+            "for" => Keyword::For,
+            "do" => Keyword::Do,
+            "break" => Keyword::Break,
+            "goto" => Keyword::GoTo,
             _ => None?,
         })
     }
@@ -50,11 +62,17 @@ impl ToString for Keyword {
             Keyword::Global => "global",
             Keyword::Return => "return",
             Keyword::If => "if",
+            Keyword::ElseIf => "elif",
+            Keyword::Else => "else",
             Keyword::Then => "then",
             Keyword::True => "true",
             Keyword::False => "false",
             Keyword::Nil => "nil",
             Keyword::Not => "not",
+            Keyword::For => "for",
+            Keyword::Do => "do",
+            Keyword::Break => "break",
+            Keyword::GoTo => "goto",
         }
         .to_string()
     }

src/vm.rs

@@ -150,8 +150,16 @@ pub enum Instruction {
     And(u16, u16, u16),
     /// PC += sAx
     Jmp(i32),
+    /// PC = Ax
+    GoTo(u32),
     /// if (R(B) <=> C) then R(A) := R(B) else PC++
     Test(u16, u16, u16),
+    /// if R(C) >= 0
+    ///     then if R(A) < R(B)
+    ///         PC ++
+    ///     otherwise if R(A) > R(B)
+    ///         PC ++
+    ForTest(u16, u16, u16),
     /// [func] [params.len()] [ret_regs.len()]
     /// R(A), ... R(A+C-2) := R(A)(R(A+1), ... R(A+B-1))
     Call(u16, u16, u16),
@@ -184,7 +192,11 @@ impl Debug for Instruction {
             Instruction::SetList(arg0, arg1) => write!(f, "SETLIST {} {}", arg0, arg1),
             Instruction::NewTable(arg0) => write!(f, "NEWTABLE {}", arg0),
             Instruction::Jmp(arg0) => write!(f, "JMP {}", arg0),
+            Instruction::GoTo(arg0) => write!(f, "GOTO {}", arg0),
             Instruction::Test(arg0, arg1, arg2) => write!(f, "TEST {} {} {}", arg0, arg1, arg2),
+            Instruction::ForTest(arg0, arg1, arg2) => {
+                write!(f, "FORTEST {} {} {}", arg0, arg1, arg2)
+            }
             Instruction::Call(arg0, arg1, arg2) => write!(f, "CALL {} {} {}", arg0, arg1, arg2),
             Instruction::Close(arg0) => write!(f, "CLOSE {}", arg0),
             Instruction::Closure(arg0, arg1) => write!(f, "CLOSURE {} {}", arg0, arg1),
@@ -881,6 +893,10 @@ impl ClosureRunner {
                 Instruction::Jmp(b) => {
                     self.program_counter = (self.program_counter as i32 + *b) as usize
                 }
+                Instruction::GoTo(a) => {
+                    self.program_counter = *a as usize;
+                    return Ok(None);
+                }
                 Instruction::Test(a, b, c) => {
                     let is_true = match self
                         .stack
@@ -888,7 +904,7 @@ impl ClosureRunner {
                         .map(|v| v.borrow().clone())
                         .unwrap_or(Value::Nil)
                     {
-                        Value::Float(val) => (val.lua_number().0 as u16) == *c,
+                        Value::Boolean(val) => (val.0 as u16) == *c,
                         _ => false,
                     };
                     if is_true {
@@ -898,6 +914,27 @@ impl ClosureRunner {
                         self.program_counter += 1;
                     }
                 }
+                Instruction::ForTest(counter_reg, end_reg, step_reg) => {
+                    let step_gte_0 = match self.get_stack(*step_reg) {
+                        StackValue::Value(value) => Value::Integer(LuaInteger(0)).lt(&value),
+                    };
+                    let counter = match self.get_stack(*counter_reg) {
+                        StackValue::Value(value) => value,
+                    };
+                    let end = match self.get_stack(*end_reg) {
+                        StackValue::Value(value) => value,
+                    };
+
+                    if step_gte_0?.is_truthy() {
+                        if !end.lt(&counter)?.is_truthy() {
+                            self.program_counter += 1;
+                        }
+                    } else {
+                        if !counter.lt(&end)?.is_truthy() {
+                            self.program_counter += 1;
+                        }
+                    }
+                }
                 Instruction::Call(func_reg, param_len, ret_len) => {
                     let param_start_func_reg = *func_reg;