diff --git a/compiler/noirc_frontend/src/ast/statement.rs b/compiler/noirc_frontend/src/ast/statement.rs index 9fe78f40c59..7d19a079203 100644 --- a/compiler/noirc_frontend/src/ast/statement.rs +++ b/compiler/noirc_frontend/src/ast/statement.rs @@ -1,9 +1,14 @@ use std::fmt::Display; +use std::sync::atomic::{AtomicU32, Ordering}; use crate::lexer::token::SpannedToken; use crate::parser::{ParserError, ParserErrorReason}; use crate::token::Token; -use crate::{Expression, ExpressionKind, IndexExpression, MemberAccessExpression, UnresolvedType}; +use crate::{ + BlockExpression, Expression, ExpressionKind, IndexExpression, MemberAccessExpression, + MethodCallExpression, UnresolvedType, +}; +use acvm::FieldElement; use iter_extended::vecmap; use noirc_errors::{Span, Spanned}; @@ -478,12 +483,123 @@ impl LValue { } } +#[derive(Debug, PartialEq, Eq, Clone)] +pub enum ForRange { + Range(/*start:*/ Expression, /*end:*/ Expression), + Array(Expression), +} + +impl ForRange { + /// Create a 'for' expression taking care of desugaring a 'for e in array' loop + /// into the following if needed: + /// + /// { + /// let fresh1 = array; + /// for fresh2 in 0 .. std::array::len(fresh1) { + /// let elem = fresh1[fresh2]; + /// ... + /// } + /// } + pub(crate) fn into_for( + self, + identifier: Ident, + block: Expression, + for_loop_span: Span, + ) -> StatementKind { + /// Counter used to generate unique names when desugaring + /// code in the parser requires the creation of fresh variables. + /// The parser is stateless so this is a static global instead. + static UNIQUE_NAME_COUNTER: AtomicU32 = AtomicU32::new(0); + + match self { + ForRange::Range(..) => { + unreachable!() + } + ForRange::Array(array) => { + let array_span = array.span; + let start_range = ExpressionKind::integer(FieldElement::zero()); + let start_range = Expression::new(start_range, array_span); + + let next_unique_id = UNIQUE_NAME_COUNTER.fetch_add(1, Ordering::Relaxed); + let array_name = format!("$i{next_unique_id}"); + let array_span = array.span; + let array_ident = Ident::new(array_name, array_span); + + // let fresh1 = array; + let let_array = Statement { + kind: StatementKind::Let(LetStatement { + pattern: Pattern::Identifier(array_ident.clone()), + r#type: UnresolvedType::unspecified(), + expression: array, + }), + span: array_span, + }; + + // array.len() + let segments = vec![array_ident]; + let array_ident = + ExpressionKind::Variable(Path { segments, kind: PathKind::Plain }); + + let end_range = ExpressionKind::MethodCall(Box::new(MethodCallExpression { + object: Expression::new(array_ident.clone(), array_span), + method_name: Ident::new("len".to_string(), array_span), + arguments: vec![], + })); + let end_range = Expression::new(end_range, array_span); + + let next_unique_id = UNIQUE_NAME_COUNTER.fetch_add(1, Ordering::Relaxed); + let index_name = format!("$i{next_unique_id}"); + let fresh_identifier = Ident::new(index_name.clone(), array_span); + + // array[i] + let segments = vec![Ident::new(index_name, array_span)]; + let index_ident = + ExpressionKind::Variable(Path { segments, kind: PathKind::Plain }); + + let loop_element = ExpressionKind::Index(Box::new(IndexExpression { + collection: Expression::new(array_ident, array_span), + index: Expression::new(index_ident, array_span), + })); + + // let elem = array[i]; + let let_elem = Statement { + kind: StatementKind::Let(LetStatement { + pattern: Pattern::Identifier(identifier), + r#type: UnresolvedType::unspecified(), + expression: Expression::new(loop_element, array_span), + }), + span: array_span, + }; + + let block_span = block.span; + let new_block = BlockExpression(vec![ + let_elem, + Statement { kind: StatementKind::Expression(block), span: block_span }, + ]); + let new_block = Expression::new(ExpressionKind::Block(new_block), block_span); + let for_loop = Statement { + kind: StatementKind::For(ForLoopStatement { + identifier: fresh_identifier, + range: ForRange::Range(start_range, end_range), + block: new_block, + span: for_loop_span, + }), + span: for_loop_span, + }; + + let block = ExpressionKind::Block(BlockExpression(vec![let_array, for_loop])); + StatementKind::Expression(Expression::new(block, for_loop_span)) + } + } + } +} + #[derive(Debug, PartialEq, Eq, Clone)] pub struct ForLoopStatement { pub identifier: Ident, - pub start_range: Expression, - pub end_range: Expression, + pub range: ForRange, pub block: Expression, + pub span: Span, } impl Display for StatementKind { @@ -575,10 +691,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) } } diff --git a/compiler/noirc_frontend/src/hir/resolution/resolver.rs b/compiler/noirc_frontend/src/hir/resolution/resolver.rs index b3796418ab3..72597ae97ba 100644 --- a/compiler/noirc_frontend/src/hir/resolution/resolver.rs +++ b/compiler/noirc_frontend/src/hir/resolution/resolver.rs @@ -36,10 +36,11 @@ use crate::{ StatementKind, }; use crate::{ - ArrayLiteral, ContractFunctionType, Distinctness, FunctionVisibility, Generics, LValue, - NoirStruct, NoirTypeAlias, Path, PathKind, Pattern, Shared, StructType, Type, TypeAliasType, - TypeBinding, TypeVariable, UnaryOp, UnresolvedGenerics, UnresolvedTraitConstraint, - UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression, Visibility, ERROR_IDENT, + ArrayLiteral, ContractFunctionType, Distinctness, ForRange, FunctionVisibility, Generics, + LValue, NoirStruct, NoirTypeAlias, Path, PathKind, Pattern, Shared, StructType, Type, + TypeAliasType, TypeBinding, TypeVariable, UnaryOp, UnresolvedGenerics, + UnresolvedTraitConstraint, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression, + Visibility, ERROR_IDENT, }; use fm::FileId; use iter_extended::vecmap; @@ -1007,23 +1008,37 @@ impl<'a> Resolver<'a> { 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, } diff --git a/compiler/noirc_frontend/src/parser/mod.rs b/compiler/noirc_frontend/src/parser/mod.rs index 4b76e2020b5..a6c631895cd 100644 --- a/compiler/noirc_frontend/src/parser/mod.rs +++ b/compiler/noirc_frontend/src/parser/mod.rs @@ -11,17 +11,13 @@ mod labels; #[allow(clippy::module_inception)] mod parser; -use std::sync::atomic::{AtomicU32, Ordering}; - use crate::token::{Keyword, Token}; use crate::{ast::ImportStatement, Expression, NoirStruct}; use crate::{ - BlockExpression, ExpressionKind, ForLoopStatement, Ident, IndexExpression, LetStatement, - MethodCallExpression, NoirFunction, NoirTrait, NoirTraitImpl, NoirTypeAlias, Path, PathKind, - Pattern, Recoverable, Statement, StatementKind, TypeImpl, UnresolvedType, UseTree, + Ident, LetStatement, NoirFunction, NoirTrait, NoirTraitImpl, NoirTypeAlias, Recoverable, + StatementKind, TypeImpl, UseTree, }; -use acvm::FieldElement; use chumsky::prelude::*; use chumsky::primitive::Container; pub use errors::ParserError; @@ -29,11 +25,6 @@ pub use errors::ParserErrorReason; use noirc_errors::Span; pub use parser::parse_program; -/// Counter used to generate unique names when desugaring -/// code in the parser requires the creation of fresh variables. -/// The parser is stateless so this is a static global instead. -static UNIQUE_NAME_COUNTER: AtomicU32 = AtomicU32::new(0); - #[derive(Debug, Clone)] pub(crate) enum TopLevelStatement { Function(NoirFunction), @@ -478,106 +469,6 @@ impl Precedence { } } -enum ForRange { - Range(/*start:*/ Expression, /*end:*/ Expression), - Array(Expression), -} - -impl ForRange { - /// Create a 'for' expression taking care of desugaring a 'for e in array' loop - /// into the following if needed: - /// - /// { - /// let fresh1 = array; - /// for fresh2 in 0 .. std::array::len(fresh1) { - /// let elem = fresh1[fresh2]; - /// ... - /// } - /// } - 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::Array(array) => { - let array_span = array.span; - let start_range = ExpressionKind::integer(FieldElement::zero()); - let start_range = Expression::new(start_range, array_span); - - let next_unique_id = UNIQUE_NAME_COUNTER.fetch_add(1, Ordering::Relaxed); - let array_name = format!("$i{next_unique_id}"); - let array_span = array.span; - let array_ident = Ident::new(array_name, array_span); - - // let fresh1 = array; - let let_array = Statement { - kind: StatementKind::Let(LetStatement { - pattern: Pattern::Identifier(array_ident.clone()), - r#type: UnresolvedType::unspecified(), - expression: array, - }), - span: array_span, - }; - - // array.len() - let segments = vec![array_ident]; - let array_ident = - ExpressionKind::Variable(Path { segments, kind: PathKind::Plain }); - - let end_range = ExpressionKind::MethodCall(Box::new(MethodCallExpression { - object: Expression::new(array_ident.clone(), array_span), - method_name: Ident::new("len".to_string(), array_span), - arguments: vec![], - })); - let end_range = Expression::new(end_range, array_span); - - let next_unique_id = UNIQUE_NAME_COUNTER.fetch_add(1, Ordering::Relaxed); - let index_name = format!("$i{next_unique_id}"); - let fresh_identifier = Ident::new(index_name.clone(), array_span); - - // array[i] - let segments = vec![Ident::new(index_name, array_span)]; - let index_ident = - ExpressionKind::Variable(Path { segments, kind: PathKind::Plain }); - - let loop_element = ExpressionKind::Index(Box::new(IndexExpression { - collection: Expression::new(array_ident, array_span), - index: Expression::new(index_ident, array_span), - })); - - // let elem = array[i]; - let let_elem = Statement { - kind: StatementKind::Let(LetStatement { - pattern: Pattern::Identifier(identifier), - r#type: UnresolvedType::unspecified(), - expression: Expression::new(loop_element, array_span), - }), - span: array_span, - }; - - let block_span = block.span; - let new_block = BlockExpression(vec![ - let_elem, - Statement { kind: StatementKind::Expression(block), span: block_span }, - ]); - let new_block = Expression::new(ExpressionKind::Block(new_block), block_span); - let for_loop = Statement { - kind: StatementKind::For(ForLoopStatement { - identifier: fresh_identifier, - start_range, - end_range, - block: new_block, - }), - span: for_loop_span, - }; - - let block = ExpressionKind::Block(BlockExpression(vec![let_array, for_loop])); - StatementKind::Expression(Expression::new(block, for_loop_span)) - } - } - } -} - impl std::fmt::Display for TopLevelStatement { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { diff --git a/compiler/noirc_frontend/src/parser/parser.rs b/compiler/noirc_frontend/src/parser/parser.rs index cee824a51c9..6d934421783 100644 --- a/compiler/noirc_frontend/src/parser/parser.rs +++ b/compiler/noirc_frontend/src/parser/parser.rs @@ -26,8 +26,8 @@ use super::{ foldl_with_span, labels::ParsingRuleLabel, parameter_name_recovery, parameter_recovery, parenthesized, then_commit, then_commit_ignore, top_level_statement_recovery, ExprParser, - ForRange, NoirParser, ParsedModule, ParsedSubModule, ParserError, ParserErrorReason, - Precedence, TopLevelStatement, + NoirParser, ParsedModule, ParsedSubModule, ParserError, ParserErrorReason, Precedence, + TopLevelStatement, }; use super::{spanned, Item, ItemKind}; use crate::ast::{ @@ -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, ForRange, 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. diff --git a/tooling/nargo_fmt/src/visitor/stmt.rs b/tooling/nargo_fmt/src/visitor/stmt.rs index 32f1126cb9d..b6dd67323fa 100644 --- a/tooling/nargo_fmt/src/visitor/stmt.rs +++ b/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::{ + ConstrainKind, ConstrainStatement, ExpressionKind, ForRange, 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!(), diff --git a/tooling/nargo_fmt/tests/expected/for.nr b/tooling/nargo_fmt/tests/expected/for.nr new file mode 100644 index 00000000000..98dff672bef --- /dev/null +++ b/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; + } +} diff --git a/tooling/nargo_fmt/tests/input/for.nr b/tooling/nargo_fmt/tests/input/for.nr new file mode 100644 index 00000000000..99b796df820 --- /dev/null +++ b/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; + } +}