✨ For loops

examples/prime.basic

@@ -0,0 +1,15 @@
+READ(input, cell1, 0)
+
+IF input % 2 == 0 THEN
+    prime = false
+ELSE
+    FOR k = 3 TO SQRT(input) STEP 2
+        IF input % k == 0 THEN
+            prime = false
+            GOTO exit
+        END IF
+    NEXT k
+    prime = true
+END IF
+exit:
+    WRITE(prime, cell1, 1)

examples/procedural-waves.basic

@@ -61,8 +61,7 @@ main:
         GOTO main_wait
 
 spawn_tank:
-    LET spawned = 0
-    spawn_tank_loop:
+    FOR spawned = 1 TO tank_units
         LET roll = rand(progression)
         IF roll >= 3 THEN
             IF roll >= 4 THEN
@@ -92,15 +91,11 @@ spawn_tank:
             spawnx = 10
         END IF
 
-        spawned = spawned + 1
-        IF spawned < tank_units THEN
-            GOTO spawn_tank_loop
-        END IF
-        GOTO spawn_tank_end
+    NEXT spawned
+    GOTO spawn_tank_end
 
 spawn_mech:
-    LET spawned = 0
-    spawn_mech_loop:
+    FOR spawned = 1 TO mech_units
         LET roll = rand(progression)
         IF roll >= 3 THEN
             IF roll >= 4 THEN
@@ -129,15 +124,11 @@ spawn_mech:
             spawnx = 10
         END IF
 
-        spawned = spawned + 1
-        IF spawned < mech_units THEN
-            GOTO spawn_mech_loop
-        END IF
-        GOTO spawn_mech_end
+    NEXT spawned
+    GOTO spawn_mech_end
 
 spawn_air:
-    LET spawned = 0
-    spawn_air_loop:
+    FOR spawned = 1 TO air_units
         LET roll = rand(progression)
         IF roll >= 3 THEN
             IF roll >= 4 THEN
@@ -166,8 +157,5 @@ spawn_air:
             spawnx = 10
         END IF
 
-        spawned = spawned + 1
-        IF spawned < air_units THEN
-            GOTO spawn_air_loop
-        END IF
-        GOTO spawn_air_end
+    NEXT spawned
+    GOTO spawn_air_end

src/compile/mod.rs

@@ -54,6 +54,7 @@ impl MindustryOperation {
         match self {
             Self::JumpIf(_label, _operator, lhs, rhs) => Box::new([lhs, rhs]),
             Self::Operator(_target, _operator, lhs, rhs) => Box::new([lhs, rhs]),
+            Self::UnaryOperator(_target, _operator, value) => Box::new([value]),
             Self::Set(_target, value) => Box::new([value]),
             Self::Generic(_name, operands) => {
                 operands.iter().collect::<Vec<_>>().into_boxed_slice()
@@ -69,6 +70,7 @@ impl MindustryOperation {
         match self {
             Self::JumpIf(_label, _operator, lhs, rhs) => vec![lhs, rhs],
             Self::Operator(_target, _operator, lhs, rhs) => vec![lhs, rhs],
+            Self::UnaryOperator(_target, _operator, value) => vec![value],
             Self::Set(_target, value) => vec![value],
             Self::Generic(_name, operands) => operands.iter_mut().collect::<Vec<_>>(),
             Self::GenericMut(_name, _out_name, operands) => operands.iter_mut().collect::<Vec<_>>(),
@@ -184,12 +186,13 @@ fn translate_expression(
             res
         }
         BasicAstExpression::Unary(op, value) => {
-            let mut res = translate_expression(value.as_ref(), namer, target_name.clone());
+            let tmp_name = namer.temporary();
+            let mut res = translate_expression(value.as_ref(), namer, tmp_name.clone());
 
             res.push(MindustryOperation::UnaryOperator(
                 target_name.clone(),
                 *op,
-                Operand::Variable(target_name),
+                Operand::Variable(tmp_name),
             ));
 
             res
@@ -257,6 +260,40 @@ pub fn translate_ast(
                     res.push(MindustryOperation::JumpLabel(end_label));
                 }
             }
+            Instr::For {
+                variable,
+                start,
+                end,
+                step,
+                instructions,
+            } => {
+                let start_label = namer.label("start");
+                let end_name = namer.temporary();
+                let step_name = namer.temporary();
+
+                // Initialization: evaluate `start`, `end` and `step`
+                res.append(&mut translate_expression(start, namer, variable.clone()));
+                res.append(&mut translate_expression(end, namer, end_name.clone()));
+                res.append(&mut translate_expression(step, namer, step_name.clone()));
+
+                // Body
+                res.push(MindustryOperation::JumpLabel(start_label.clone()));
+                res.append(&mut translate_ast(instructions, namer, config));
+
+                // Loop condition: increment variable and jump
+                res.push(MindustryOperation::Operator(
+                    variable.clone(),
+                    Operator::Add,
+                    Operand::Variable(variable.clone()),
+                    Operand::Variable(step_name),
+                ));
+                res.push(MindustryOperation::JumpIf(
+                    start_label,
+                    Operator::Lte,
+                    Operand::Variable(variable.clone()),
+                    Operand::Variable(end_name),
+                ));
+            }
             Instr::Print(expressions) => {
                 for expression in expressions {
                     let tmp_name = namer.temporary();

src/compile/optimize/constant.rs

@@ -21,16 +21,14 @@ pub fn optimize_constant(program: MindustryProgram) -> MindustryProgram {
                 Operand::Variable(name) if tmp_regex.is_match(name) => Some(name),
                 _ => None,
             })
-            // PERF: check when it would be better to deduplicate operands
-            // .collect::<HashSet<_>>()
-            // .into_iter()
             .filter_map(|name| {
                 lookbehind(instructions, use_index, |instr| {
                     match instr {
                         MindustryOperation::Set(set_name, value) if set_name == name => {
                             Lookaround::Stop((name.clone(), value.clone()))
                         }
-                        MindustryOperation::Operator(op_name, _op, _lhs, _rhs)
+                        MindustryOperation::Operator(op_name, _, _, _)
+                        | MindustryOperation::UnaryOperator(op_name, _, _)
                             if op_name == name =>
                         {
                             Lookaround::Abort
@@ -55,7 +53,13 @@ pub fn optimize_constant(program: MindustryProgram) -> MindustryProgram {
                             // but this usecase isn't needed yet.
                             Lookaround::Abort
                         }
-                        _ => Lookaround::Continue,
+                        other => {
+                            if other.mutates(name) {
+                                Lookaround::Abort
+                            } else {
+                                Lookaround::Continue
+                            }
+                        }
                     }
                 })
                 .map(|(index, (name, value))| (name, value, index))

src/parse/ast.rs

@@ -20,6 +20,13 @@ pub enum BasicAstInstruction {
     IfThenElse(BasicAstExpression, BasicAstBlock, BasicAstBlock),
     Print(Vec<(BasicAstExpression, bool)>),
     CallBuiltin(String, Vec<BasicAstExpression>),
+    For {
+        variable: String,
+        start: BasicAstExpression,
+        end: BasicAstExpression,
+        step: BasicAstExpression,
+        instructions: BasicAstBlock,
+    },
 }
 
 #[derive(Clone, Debug, PartialEq, Default)]
@@ -34,12 +41,17 @@ impl BasicAstBlock {
         }
     }
 }
-
 pub fn build_ast(tokens: &[BasicToken], config: &Config) -> Result<BasicAstBlock, ParseError> {
     enum Context {
         Main,
         If(BasicAstExpression),
         IfElse(BasicAstExpression, BasicAstBlock),
+        For(
+            String,
+            BasicAstExpression,
+            BasicAstExpression,
+            BasicAstExpression,
+        ),
     }
 
     let mut tokens = Cursor::from(tokens);
@@ -63,14 +75,7 @@ pub fn build_ast(tokens: &[BasicToken], config: &Config) -> Result<BasicAstBlock
             [BasicToken::NewLine, ..] => {
                 tokens.take(1);
             }
-            [BasicToken::Name(variable_name), BasicToken::Assign, ..] => {
-                tokens.take(2);
-                let expression = parse_expression(&mut tokens)?;
-                instructions.push(BasicAstInstruction::Assign(
-                    variable_name.clone(),
-                    expression,
-                ));
-            }
+            // == If-then-else ==
             [BasicToken::If, ..] => {
                 tokens.take(1);
                 let then_index = find_token_index(&tokens, BasicToken::Then)?;
@@ -129,6 +134,62 @@ pub fn build_ast(tokens: &[BasicToken], config: &Config) -> Result<BasicAstBlock
                     }
                 }
             }
+            // == For loops ==
+            [BasicToken::For, BasicToken::Name(variable), BasicToken::Assign, ..] => {
+                tokens.take(3);
+
+                let start = parse_expression(&mut tokens)?;
+
+                expect_next_token(&mut tokens, &BasicToken::To)?;
+                tokens.take(1);
+
+                let end = parse_expression(&mut tokens)?;
+
+                let step = if let Some(BasicToken::Step) = tokens.get(0) {
+                    tokens.take(1);
+
+                    parse_expression(&mut tokens)?
+                } else {
+                    BasicAstExpression::Integer(1)
+                };
+
+                expect_next_token(&mut tokens, &BasicToken::NewLine)?;
+
+                context_stack.push((Vec::new(), Context::For(variable.clone(), start, end, step)));
+            }
+            [BasicToken::Next, BasicToken::Name(variable), ..] => match context_stack.pop() {
+                Some((instructions, Context::For(expected_variable, start, end, step))) => {
+                    tokens.take(2);
+
+                    let Some((ref mut parent_instructions, _)) = context_stack.last_mut() else {
+                        unreachable!("Context::For not wrapped in another context");
+                    };
+
+                    if *variable != expected_variable {
+                        return Err(ParseError::WrongForVariable(
+                            expected_variable,
+                            variable.clone(),
+                        ));
+                    }
+
+                    parent_instructions.push(BasicAstInstruction::For {
+                        variable: expected_variable,
+                        start,
+                        end,
+                        step,
+                        instructions: BasicAstBlock::new(instructions),
+                    });
+                }
+                Some((_instructions, _context)) => {
+                    eprintln!("NEXT outside of loop");
+                    return Err(ParseError::UnexpectedToken(BasicToken::Next));
+                }
+                None => {
+                    unreachable!("Empty context stack");
+                }
+            },
+
+            // == Goto ==
             [BasicToken::Goto, BasicToken::Integer(label), ..] => {
                 tokens.take(2);
                 instructions.push(BasicAstInstruction::Jump(label.to_string()));
@@ -137,6 +198,15 @@ pub fn build_ast(tokens: &[BasicToken], config: &Config) -> Result<BasicAstBlock
                 tokens.take(2);
                 instructions.push(BasicAstInstruction::Jump(label.clone()));
             }
+            // == Misc ==
+            [BasicToken::Name(variable_name), BasicToken::Assign, ..] => {
+                tokens.take(2);
+                let expression = parse_expression(&mut tokens)?;
+                instructions.push(BasicAstInstruction::Assign(
+                    variable_name.clone(),
+                    expression,
+                ));
+            }
             [BasicToken::Print, ..] => {
                 tokens.take(1);
 
@@ -232,6 +302,9 @@ pub fn build_ast(tokens: &[BasicToken], config: &Config) -> Result<BasicAstBlock
             Context::If(_) | Context::IfElse(_, _) => {
                 return Err(ParseError::MissingToken(BasicToken::EndIf));
             }
+            Context::For(_, _, _, _) => {
+                return Err(ParseError::MissingToken(BasicToken::Next));
+            }
             Context::Main => {
                 unreachable!("There cannot be another context below the main context");
             }
@@ -310,15 +383,8 @@ pub(crate) fn parse_expression(
                     }
                 }
 
-                match tokens.take(1) {
-                    [BasicToken::CloseParen] => {}
-                    [other] => {
-                        return Err(ParseError::UnexpectedToken(other.clone()));
-                    }
-                    _ => {
-                        return Err(ParseError::MissingToken(BasicToken::CloseParen));
-                    }
-                }
+                expect_next_token(tokens, &BasicToken::CloseParen)?;
+                tokens.take(1);
 
                 if let Ok(unary_operator) = UnaryOperator::try_from(fn_name_lowercase.as_str()) {
                     if arguments.len() != 1 {
@@ -415,3 +481,14 @@ pub(crate) fn parse_expression(
 
     Ok(res)
 }
+
+fn expect_next_token(
+    tokens: &Cursor<'_, BasicToken>,
+    expected: &BasicToken,
+) -> Result<(), ParseError> {
+    match tokens.get(0) {
+        Some(token) if token == expected => Ok(()),
+        Some(token) => Err(ParseError::UnexpectedToken(token.clone())),
+        None => Err(ParseError::MissingToken(expected.clone())),
+    }
+}

src/parse/mod.rs

@@ -15,4 +15,5 @@ pub enum ParseError {
     InvalidArgumentCount(String, usize, usize),
     ExpectedVariable,
     ExpectedOperand,
+    WrongForVariable(String, String),
 }

src/parse/test.rs

@@ -67,6 +67,46 @@ fn test_tokenize_basic() {
     );
 }
 
+#[test]
+fn test_parse_for() {
+    assert_eq!(
+        tokenize("FOR x = 0 TO y\nPRINT x\nNEXT x").unwrap(),
+        vec![
+            BasicToken::NewLine,
+            BasicToken::For,
+            BasicToken::Name(String::from("x")),
+            BasicToken::Assign,
+            BasicToken::Integer(0),
+            BasicToken::To,
+            BasicToken::Name(String::from("y")),
+            BasicToken::NewLine,
+            BasicToken::Print,
+            BasicToken::Name(String::from("x")),
+            BasicToken::NewLine,
+            BasicToken::Next,
+            BasicToken::Name(String::from("x")),
+        ]
+    );
+
+    assert_eq!(
+        build_ast(
+            &tokenize("FOR x = 0 TO y\nPRINT x\nNEXT x").unwrap(),
+            &Default::default()
+        )
+        .unwrap(),
+        BasicAstBlock::new([BasicAstInstruction::For {
+            variable: String::from("x"),
+            start: BasicAstExpression::Integer(0),
+            end: BasicAstExpression::Variable(String::from("y")),
+            step: BasicAstExpression::Integer(1),
+            instructions: BasicAstBlock::new([BasicAstInstruction::Print(vec![(
+                BasicAstExpression::Variable(String::from("x")),
+                false
+            )]),])
+        }])
+    );
+}
+
 #[test]
 fn test_operator_precedence() {
     fn test_parse<const N: usize>(list: [BasicToken; N]) -> BasicAstExpression {

src/parse/tokenize.rs

@@ -11,6 +11,10 @@ pub enum BasicToken {
     Else,
     EndIf,
     Goto,
+    For,
+    To,
+    Step,
+    Next,
     LabelEnd,
     OpenParen,
     CloseParen,
@@ -70,7 +74,8 @@ pub fn tokenize(raw: &str) -> Result<Vec<BasicToken>, ParseError> {
     let mut res = Vec::new();
     let match_let = Regex::new(r"(?i)^let").unwrap();
     let match_jump = Regex::new(r"(?i)^go\s*to").unwrap();
-    let match_word = Regex::new(r"(?i)^(?:if|then|else|end\s?if|print)(?:\s|$)").unwrap();
+    let match_word =
+        Regex::new(r"(?i)^(?:if|then|else|end\s?if|print|for|to|step|next)(?:\s|$)").unwrap();
     let match_space = Regex::new(r"^\s+").unwrap();
     let match_variable = Regex::new(r"^@?[a-zA-Z_][a-zA-Z_0-9]*").unwrap();
     let match_float = Regex::new(r"^[0-9]*\.[0-9]+").unwrap();
@@ -102,6 +107,10 @@ pub fn tokenize(raw: &str) -> Result<Vec<BasicToken>, ParseError> {
                     "else" => BasicToken::Else,
                     "end if" | "endif" => BasicToken::EndIf,
                     "print" => BasicToken::Print,
+                    "for" => BasicToken::For,
+                    "to" => BasicToken::To,
+                    "step" => BasicToken::Step,
+                    "next" => BasicToken::Next,
                     _ => unreachable!("{}", word),
                 }),
                 match_variable(name) => (BasicToken::Name(name.to_string())),