chore: format for stmt

compiler/noirc_frontend/src/ast/statement.rs

@@ -1,7 +1,7 @@
 use std::fmt::Display;
 
 use crate::lexer::token::SpannedToken;
-use crate::parser::{ParserError, ParserErrorReason};
+use crate::parser::{ForRange, ParserError, ParserErrorReason};
 use crate::token::Token;
 use crate::{Expression, ExpressionKind, IndexExpression, MemberAccessExpression, UnresolvedType};
 use iter_extended::vecmap;
@@ -480,6 +480,14 @@ impl LValue {
 
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct ForLoopStatement {
+    pub identifier: Ident,
+    pub range: ForRange,
+    pub block: Expression,
+    pub span: Span,
+}
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct ForRangeLoopStatement {
     pub identifier: Ident,
     pub start_range: Expression,
     pub end_range: Expression,
@@ -575,10 +583,11 @@ impl Display for Pattern {
 
 impl Display for ForLoopStatement {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(
-            f,
-            "for {} in {} .. {} {}",
-            self.identifier, self.start_range, self.end_range, self.block
-        )
+        let range = match &self.range {
+            ForRange::Range(start, end) => format!("{start}..{end}"),
+            ForRange::Array(expr) => expr.to_string(),
+        };
+
+        write!(f, "for {} in {range} {}", self.identifier, self.block)
     }
 }

compiler/noirc_frontend/src/hir/resolution/resolver.rs

@@ -19,6 +19,7 @@
 };
 
 use crate::hir_def::traits::{Trait, TraitConstraint};
+use crate::parser::ForRange;
 use crate::token::FunctionAttribute;
 use regex::Regex;
 use std::collections::{BTreeMap, HashSet};
@@ -89,7 +90,7 @@
 
     /// True if the current module is a contract.
     /// This is usually determined by self.path_resolver.module_id(), but it can
     /// be overriden for impls. Impls are an odd case since the methods within resolve
     /// as if they're in the parent module, but should be placed in a child module.
     /// Since they should be within a child module, in_contract is manually set to false
     /// for these so we can still resolve them in the parent module without them being in a contract.
@@ -1007,23 +1008,37 @@
                 HirStatement::Assign(stmt)
             }
             StatementKind::For(for_loop) => {
-                let start_range = self.resolve_expression(for_loop.start_range);
-                let end_range = self.resolve_expression(for_loop.end_range);
-                let (identifier, block) = (for_loop.identifier, for_loop.block);
-
-                // TODO: For loop variables are currently mutable by default since we haven't
-                //       yet implemented syntax for them to be optionally mutable.
-                let (identifier, block) = self.in_new_scope(|this| {
-                    let decl = this.add_variable_decl(
-                        identifier,
-                        false,
-                        true,
-                        DefinitionKind::Local(None),
-                    );
-                    (decl, this.resolve_expression(block))
-                });
+                match for_loop.range {
+                    ForRange::Range(start_range, end_range) => {
+                        let start_range = self.resolve_expression(start_range);
+                        let end_range = self.resolve_expression(end_range);
+                        let (identifier, block) = (for_loop.identifier, for_loop.block);
+
+                        // TODO: For loop variables are currently mutable by default since we haven't
+                        //       yet implemented syntax for them to be optionally mutable.
+                        let (identifier, block) = self.in_new_scope(|this| {
+                            let decl = this.add_variable_decl(
+                                identifier,
+                                false,
+                                true,
+                                DefinitionKind::Local(None),
+                            );
+                            (decl, this.resolve_expression(block))
+                        });
 
-                HirStatement::For(HirForStatement { start_range, end_range, block, identifier })
+                        HirStatement::For(HirForStatement {
+                            start_range,
+                            end_range,
+                            block,
+                            identifier,
+                        })
+                    }
+                    range @ ForRange::Array(_) => {
+                        let for_stmt =
+                            range.into_for(for_loop.identifier, for_loop.block, for_loop.span);
+                        self.resolve_stmt(for_stmt)
+                    }
+                }
             }
             StatementKind::Error => HirStatement::Error,
         }

compiler/noirc_frontend/src/parser/mod.rs

@@ -478,7 +478,8 @@ impl Precedence {
     }
 }
 
-enum ForRange {
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum ForRange {
     Range(/*start:*/ Expression, /*end:*/ Expression),
     Array(Expression),
 }
@@ -494,10 +495,15 @@ impl ForRange {
     ///         ...
     ///     }
     /// }
-    fn into_for(self, identifier: Ident, block: Expression, for_loop_span: Span) -> StatementKind {
+    pub(crate) fn into_for(
+        self,
+        identifier: Ident,
+        block: Expression,
+        for_loop_span: Span,
+    ) -> StatementKind {
         match self {
-            ForRange::Range(start_range, end_range) => {
-                StatementKind::For(ForLoopStatement { identifier, start_range, end_range, block })
+            ForRange::Range(..) => {
+                unreachable!()
             }
             ForRange::Array(array) => {
                 let array_span = array.span;
@@ -564,9 +570,9 @@ impl ForRange {
                 let for_loop = Statement {
                     kind: StatementKind::For(ForLoopStatement {
                         identifier: fresh_identifier,
-                        start_range,
-                        end_range,
+                        range: ForRange::Range(start_range, end_range),
                         block: new_block,
+                        span: for_loop_span,
                     }),
                     span: for_loop_span,
                 };

compiler/noirc_frontend/src/parser/parser.rs

@@ -38,11 +38,11 @@ use crate::parser::{force, ignore_then_commit, statement_recovery};
 use crate::token::{Attribute, Attributes, Keyword, SecondaryAttribute, Token, TokenKind};
 use crate::{
     BinaryOp, BinaryOpKind, BlockExpression, ConstrainKind, ConstrainStatement, Distinctness,
-    FunctionDefinition, FunctionReturnType, FunctionVisibility, Ident, IfExpression,
-    InfixExpression, LValue, Lambda, Literal, NoirFunction, NoirStruct, NoirTrait, NoirTraitImpl,
-    NoirTypeAlias, Path, PathKind, Pattern, Recoverable, Statement, TraitBound, TraitImplItem,
-    TraitItem, TypeImpl, UnaryOp, UnresolvedTraitConstraint, UnresolvedTypeExpression, UseTree,
-    UseTreeKind, Visibility,
+    ForLoopStatement, FunctionDefinition, FunctionReturnType, FunctionVisibility, Ident,
+    IfExpression, InfixExpression, LValue, Lambda, Literal, NoirFunction, NoirStruct, NoirTrait,
+    NoirTraitImpl, NoirTypeAlias, Path, PathKind, Pattern, Recoverable, Statement, TraitBound,
+    TraitImplItem, TraitItem, TypeImpl, UnaryOp, UnresolvedTraitConstraint,
+    UnresolvedTypeExpression, UseTree, UseTreeKind, Visibility,
 };
 
 use chumsky::prelude::*;
@@ -1482,7 +1482,9 @@ where
         .then_ignore(keyword(Keyword::In))
         .then(for_range(expr_no_constructors))
         .then(block_expr(statement))
-        .map_with_span(|((identifier, range), block), span| range.into_for(identifier, block, span))
+        .map_with_span(|((identifier, range), block), span| {
+            StatementKind::For(ForLoopStatement { identifier, range, block, span })
+        })
 }
 
 /// The 'range' of a for loop. Either an actual range `start .. end` or an array expression.

tooling/nargo_fmt/src/visitor/stmt.rs

@@ -1,6 +1,8 @@
 use std::iter::zip;
 
-use noirc_frontend::{ConstrainKind, ConstrainStatement, ExpressionKind, Statement, StatementKind};
+use noirc_frontend::{
+    parser::ForRange, ConstrainKind, ConstrainStatement, ExpressionKind, Statement, StatementKind,
+};
 
 use super::ExpressionType;
 
@@ -55,7 +57,22 @@ impl super::FmtVisitor<'_> {
 
                     self.push_rewrite(constrain, span);
                 }
-                StatementKind::For(_) | StatementKind::Assign(_) => {
+                StatementKind::For(for_stmt) => {
+                    let identifier = self.slice(for_stmt.identifier.span());
+                    let range = match for_stmt.range {
+                        ForRange::Range(start, end) => format!(
+                            "{}..{}",
+                            self.format_sub_expr(start),
+                            self.format_sub_expr(end)
+                        ),
+                        ForRange::Array(array) => self.format_sub_expr(array),
+                    };
+                    let block = self.format_sub_expr(for_stmt.block);
+
+                    let result = format!("for {identifier} in {range} {block}");
+                    self.push_rewrite(result, span);
+                }
+                StatementKind::Assign(_) => {
                     self.push_rewrite(self.slice(span).to_string(), span);
                 }
                 StatementKind::Error => unreachable!(),

tooling/nargo_fmt/tests/expected/for.nr

@@ -0,0 +1,21 @@
+fn for_stmt() {
+    for elem in self {
+        ret &= predicate(elem);
+    }
+}
+
+fn for_stmt() {
+    for i in 0..(C1 - 1) {
+        for _j in 1..(C1 - i - 1) {
+            b *= b;
+        }
+
+        z *= if b == 1 { 1 } else { c };
+
+        c *= c;
+
+        t *= if b == 1 { 1 } else { c };
+
+        b = t;
+    }
+}

tooling/nargo_fmt/tests/input/for.nr

@@ -0,0 +1,24 @@
+fn for_stmt() {
+    for elem in self {
+        ret &= predicate(elem);
+    }
+}
+
+fn for_stmt() {
+        for i in 0..(C1-1) {
+
+        for _j in 1..(C1-i-1) {
+
+            b *= b;
+
+        }
+
+        z *= if b == 1 { 1 } else { c };
+
+        c *= c;
+
+        t *= if b == 1 { 1 } else { c };
+
+        b = t;
+    }
+}