Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Inclusive for loop #6200

Merged
merged 17 commits into from
Oct 3, 2024
Merged
Show file tree
Hide file tree
Changes from 13 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 54 additions & 5 deletions compiler/noirc_frontend/src/ast/statement.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,9 @@ use iter_extended::vecmap;
use noirc_errors::{Span, Spanned};

use super::{
BlockExpression, ConstructorExpression, Expression, ExpressionKind, GenericTypeArgs,
IndexExpression, ItemVisibility, MemberAccessExpression, MethodCallExpression, UnresolvedType,
BinaryOpKind, BlockExpression, ConstructorExpression, Expression, ExpressionKind,
GenericTypeArgs, IndexExpression, InfixExpression, ItemVisibility, MemberAccessExpression,
MethodCallExpression, UnresolvedType,
};
use crate::ast::UnresolvedTypeData;
use crate::elaborator::types::SELF_TYPE_NAME;
Expand Down Expand Up @@ -770,13 +771,60 @@ impl LValue {
}
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct ForBounds(
/*start:*/ pub Expression,
/*end:*/ pub Expression,
/*inclusive:*/ pub bool,
aakoshh marked this conversation as resolved.
Show resolved Hide resolved
);

impl ForBounds {
/// Create a half-open range bounded inclusively below and exclusively above (`start..end`),
/// desugaring `start..=end` into `start..end+1` if necessary.
///
/// Returns the `start` and `end` expressions.
pub(crate) fn into_half_open(self) -> (Expression, Expression) {
match self {
ForBounds(start, end, false) => (start, end),
ForBounds(start, end, true) => {
let end_span = end.span;
let end = ExpressionKind::Infix(Box::new(InfixExpression {
lhs: end,
operator: Spanned::from(end_span, BinaryOpKind::Add),
rhs: Expression::new(
ExpressionKind::integer(FieldElement::from(1u32)),
end_span,
),
}));
let end = Expression::new(end, end_span);
(start, end)
}
}
}
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub enum ForRange {
Range(/*start:*/ Expression, /*end:*/ Expression),
Range(ForBounds),
Array(Expression),
}

impl ForRange {
/// Create a half-open range, bounded inclusively below and exclusively above.
pub fn range(start: Expression, end: Expression) -> Self {
Self::Range(ForBounds(start, end, false))
}

/// Create a range bounded inclusively below and above.
pub fn range_inclusive(start: Expression, end: Expression) -> Self {
Self::Range(ForBounds(start, end, true))
}

/// Create a range over some array.
pub fn array(value: Expression) -> Self {
Self::Array(value)
}

/// Create a 'for' expression taking care of desugaring a 'for e in array' loop
/// into the following if needed:
///
Expand Down Expand Up @@ -879,7 +927,7 @@ impl ForRange {
let for_loop = Statement {
kind: StatementKind::For(ForLoopStatement {
identifier: fresh_identifier,
range: ForRange::Range(start_range, end_range),
range: ForRange::range(start_range, end_range),
block: new_block,
span: for_loop_span,
}),
Expand Down Expand Up @@ -1009,7 +1057,8 @@ impl Display for Pattern {
impl Display for ForLoopStatement {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let range = match &self.range {
ForRange::Range(start, end) => format!("{start}..{end}"),
ForRange::Range(ForBounds(start, end, false)) => format!("{start}..{end}"),
ForRange::Range(ForBounds(start, end, true)) => format!("{start}..={end}"),
ForRange::Array(expr) => expr.to_string(),
};

Expand Down
8 changes: 4 additions & 4 deletions compiler/noirc_frontend/src/ast/visitor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@ use crate::{
};

use super::{
FunctionReturnType, GenericTypeArgs, IntegerBitSize, ItemVisibility, Pattern, Signedness,
TraitImplItemKind, TypePath, UnresolvedGenerics, UnresolvedTraitConstraint, UnresolvedType,
UnresolvedTypeData, UnresolvedTypeExpression,
ForBounds, FunctionReturnType, GenericTypeArgs, IntegerBitSize, ItemVisibility, Pattern,
Signedness, TraitImplItemKind, TypePath, UnresolvedGenerics, UnresolvedTraitConstraint,
UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
};

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
Expand Down Expand Up @@ -1192,7 +1192,7 @@ impl ForRange {

pub fn accept_children(&self, visitor: &mut impl Visitor) {
match self {
ForRange::Range(start, end) => {
ForRange::Range(ForBounds(start, end, _)) => {
start.accept(visitor);
end.accept(visitor);
}
Expand Down
2 changes: 1 addition & 1 deletion compiler/noirc_frontend/src/elaborator/statements.rs
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,7 @@ impl<'context> Elaborator<'context> {

pub(super) fn elaborate_for(&mut self, for_loop: ForLoopStatement) -> (HirStatement, Type) {
let (start, end) = match for_loop.range {
ForRange::Range(start, end) => (start, end),
ForRange::Range(bounds) => bounds.into_half_open(),
ForRange::Array(_) => {
let for_stmt =
for_loop.range.into_for(for_loop.identifier, for_loop.block, for_loop.span);
Expand Down
8 changes: 5 additions & 3 deletions compiler/noirc_frontend/src/hir/comptime/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use crate::{
ast::{
ArrayLiteral, AsTraitPath, AssignStatement, BlockExpression, CallExpression,
CastExpression, ConstrainStatement, ConstructorExpression, Expression, ExpressionKind,
ForLoopStatement, ForRange, GenericTypeArgs, IfExpression, IndexExpression,
ForBounds, ForLoopStatement, ForRange, GenericTypeArgs, IfExpression, IndexExpression,
InfixExpression, LValue, Lambda, LetStatement, Literal, MemberAccessExpression,
MethodCallExpression, Pattern, PrefixExpression, Statement, StatementKind, UnresolvedType,
UnresolvedTypeData,
Expand Down Expand Up @@ -267,6 +267,7 @@ impl<'interner> TokenPrettyPrinter<'interner> {
| Token::Dot
| Token::DoubleColon
| Token::DoubleDot
| Token::DoubleDotEqual
| Token::Caret
| Token::Pound
| Token::Pipe
Expand Down Expand Up @@ -713,10 +714,11 @@ fn remove_interned_in_statement_kind(
}),
StatementKind::For(for_loop) => StatementKind::For(ForLoopStatement {
range: match for_loop.range {
ForRange::Range(from, to) => ForRange::Range(
ForRange::Range(ForBounds(from, to, inclusive)) => ForRange::Range(ForBounds(
remove_interned_in_expression(interner, from),
remove_interned_in_expression(interner, to),
),
inclusive,
)),
ForRange::Array(expr) => {
ForRange::Array(remove_interned_in_expression(interner, expr))
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ impl HirStatement {
}),
HirStatement::For(for_stmt) => StatementKind::For(ForLoopStatement {
identifier: for_stmt.identifier.to_display_ast(interner),
range: ForRange::Range(
range: ForRange::range(
for_stmt.start_range.to_display_ast(interner),
for_stmt.end_range.to_display_ast(interner),
),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1514,7 +1514,8 @@ fn expr_as_for_range(
) -> IResult<Value> {
expr_as(interner, arguments, return_type, location, |expr| {
if let ExprValue::Statement(StatementKind::For(for_statement)) = expr {
if let ForRange::Range(from, to) = for_statement.range {
if let ForRange::Range(bounds) = for_statement.range {
let (from, to) = bounds.into_half_open();
asterite marked this conversation as resolved.
Show resolved Hide resolved
let identifier =
Value::Quoted(Rc::new(vec![Token::Ident(for_statement.identifier.0.contents)]));
let from = Value::expression(from.kind);
Expand Down
13 changes: 13 additions & 0 deletions compiler/noirc_frontend/src/hir/comptime/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,19 @@ fn for_loop() {
assert_eq!(result, Value::U8(15));
}

#[test]
fn for_loop_inclusive() {
asterite marked this conversation as resolved.
Show resolved Hide resolved
let program = "comptime fn main() -> pub u8 {
let mut x = 0;
for i in 0 ..= 6 {
x += i;
}
x
}";
let result = interpret(program);
assert_eq!(result, Value::U8(21));
}

#[test]
fn for_loop_u16() {
let program = "comptime fn main() -> pub u16 {
Expand Down
31 changes: 24 additions & 7 deletions compiler/noirc_frontend/src/lexer/lexer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,11 @@ impl<'a> Lexer<'a> {
self.peek_char() == Some(ch)
}

/// Peeks at the character two positions ahead and returns true if it is equal to the char argument
fn peek2_char_is(&mut self, ch: char) -> bool {
self.peek2_char() == Some(ch)
}

fn ampersand(&mut self) -> SpannedTokenResult {
if self.peek_char_is('&') {
// When we issue this error the first '&' will already be consumed
Expand Down Expand Up @@ -152,6 +157,7 @@ impl<'a> Lexer<'a> {
Ok(token.into_single_span(self.position))
}

/// If `single` is followed by `character` then extend it as `double`.
fn single_double_peek_token(
&mut self,
character: char,
Expand All @@ -169,13 +175,23 @@ impl<'a> Lexer<'a> {
}
}

/// Given that some tokens can contain two characters, such as <= , !=, >=
/// Glue will take the first character of the token and check if it can be glued onto the next character
/// forming a double token
/// Given that some tokens can contain two characters, such as <= , !=, >=, or even three like ..=
/// Glue will take the first character of the token and check if it can be glued onto the next character(s)
/// forming a double or triple token
///
/// Returns an error if called with a token which cannot be extended with anything.
fn glue(&mut self, prev_token: Token) -> SpannedTokenResult {
let spanned_prev_token = prev_token.clone().into_single_span(self.position);
match prev_token {
Token::Dot => self.single_double_peek_token('.', prev_token, Token::DoubleDot),
Token::Dot => {
if self.peek_char_is('.') && self.peek2_char_is('=') {
let start = self.position;
self.next_char();
self.next_char();
Ok(Token::DoubleDotEqual.into_span(start, start + 2))
} else {
self.single_double_peek_token('.', prev_token, Token::DoubleDot)
}
}
Token::Less => {
let start = self.position;
if self.peek_char_is('=') {
Expand Down Expand Up @@ -214,7 +230,7 @@ impl<'a> Lexer<'a> {
return self.parse_block_comment(start);
}

Ok(spanned_prev_token)
Ok(prev_token.into_single_span(start))
}
_ => Err(LexerErrorKind::NotADoubleChar {
span: Span::single_char(self.position),
Expand Down Expand Up @@ -704,7 +720,7 @@ mod tests {

#[test]
fn test_single_double_char() {
let input = "! != + ( ) { } [ ] | , ; : :: < <= > >= & - -> . .. % / * = == << >>";
let input = "! != + ( ) { } [ ] | , ; : :: < <= > >= & - -> . .. ..= % / * = == << >>";
asterite marked this conversation as resolved.
Show resolved Hide resolved

let expected = vec![
Token::Bang,
Expand All @@ -730,6 +746,7 @@ mod tests {
Token::Arrow,
Token::Dot,
Token::DoubleDot,
Token::DoubleDotEqual,
Token::Percent,
Token::Slash,
Token::Star,
Expand Down
6 changes: 6 additions & 0 deletions compiler/noirc_frontend/src/lexer/token.rs
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,8 @@ pub enum BorrowedToken<'input> {
Dot,
/// ..
DoubleDot,
/// ..=
DoubleDotEqual,
/// (
LeftParen,
/// )
Expand Down Expand Up @@ -190,6 +192,8 @@ pub enum Token {
Dot,
/// ..
DoubleDot,
/// ..=
DoubleDotEqual,
/// (
LeftParen,
/// )
Expand Down Expand Up @@ -279,6 +283,7 @@ pub fn token_to_borrowed_token(token: &Token) -> BorrowedToken<'_> {
Token::ShiftRight => BorrowedToken::ShiftRight,
Token::Dot => BorrowedToken::Dot,
Token::DoubleDot => BorrowedToken::DoubleDot,
Token::DoubleDotEqual => BorrowedToken::DoubleDotEqual,
Token::LeftParen => BorrowedToken::LeftParen,
Token::RightParen => BorrowedToken::RightParen,
Token::LeftBrace => BorrowedToken::LeftBrace,
Expand Down Expand Up @@ -409,6 +414,7 @@ impl fmt::Display for Token {
Token::ShiftRight => write!(f, ">>"),
Token::Dot => write!(f, "."),
Token::DoubleDot => write!(f, ".."),
Token::DoubleDotEqual => write!(f, "..="),
Token::LeftParen => write!(f, "("),
Token::RightParen => write!(f, ")"),
Token::LeftBrace => write!(f, "{{"),
Expand Down
22 changes: 17 additions & 5 deletions compiler/noirc_frontend/src/parser/parser.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1076,9 +1076,15 @@ where
{
expr_no_constructors
.clone()
.then_ignore(just(Token::DoubleDot))
.then(just(Token::DoubleDot).or(just(Token::DoubleDotEqual)))
.then(expr_no_constructors.clone())
.map(|(start, end)| ForRange::Range(start, end))
.map(|((start, dots), end)| {
if dots == Token::DoubleDotEqual {
ForRange::range_inclusive(start, end)
} else {
ForRange::range(start, end)
}
})
.or(expr_no_constructors.map(ForRange::Array))
}

Expand Down Expand Up @@ -1465,15 +1471,21 @@ mod test {
fn parse_for_loop() {
parse_all(
for_loop(expression_no_constructors(expression()), fresh_statement()),
vec!["for i in x+y..z {}", "for i in 0..100 { foo; bar }"],
vec![
"for i in x+y..z {}",
"for i in 0..100 { foo; bar }",
"for i in 90..=100 { foo; bar }",
],
);

parse_all_failing(
for_loop(expression_no_constructors(expression()), fresh_statement()),
vec![
"for 1 in x+y..z {}", // Cannot have a literal as the loop identifier
"for i in 0...100 {}", // Only '..' is supported, there are no inclusive ranges yet
"for i in 0..=100 {}", // Only '..' is supported, there are no inclusive ranges yet
"for i in 0...100 {}", // Only '..' is supported
"for i in .. {}", // Range needs needs bounds
"for i in ..=10 {}", // Range needs lower bound
"for i in 0.. {}", // Range needs upper bound
],
);
}
Expand Down
12 changes: 12 additions & 0 deletions compiler/noirc_frontend/src/tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1071,6 +1071,18 @@ fn resolve_for_expr() {
assert_no_errors(src);
}

#[test]
fn resolve_for_expr_incl() {
let src = r#"
fn main(x : u64) {
for i in 1..=20 {
let _z = x + i;
};
}
"#;
assert_no_errors(src);
}

#[test]
fn resolve_call_expr() {
let src = r#"
Expand Down
2 changes: 2 additions & 0 deletions docs/docs/noir/concepts/control_flow.md
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ for i in 0..10 {
}
```

Alternatively, `start..=end` can be used for a range that is inclusive on both ends.

The index for loops is of type `u64`.

### Break and Continue
Expand Down
2 changes: 1 addition & 1 deletion tooling/nargo_fmt/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ fn main() {
// is the root of the repository and append the crate path
let manifest_dir = match std::env::var("CARGO_MANIFEST_DIR") {
Ok(dir) => PathBuf::from(dir),
Err(_) => std::env::current_dir().unwrap().join("crates").join("nargo_cli"),
Err(_) => std::env::current_dir().unwrap().join("tooling").join("nargo_fmt"),
};
let test_dir = manifest_dir.join("tests");

Expand Down
Loading
Loading