Skip to content

Commit

Permalink
feat: Inclusive for loop (#6200)
Browse files Browse the repository at this point in the history
# Description

## Problem\*

Resolves #6171

## Summary\*

Adds support for inclusive for syntax, e.g. `for i in 0..=10`.

## Additional Context

- Added a `ForBounds` struct to handle the desugaring of `start..=end`
into `start..end+1`.
- `ForBounds` has an `inclusive` flag; I thought about adding a
`ForRange::RangeInclusive` similar to `std::ops::RangeInclusive` in
Rust, but then I thought `for` loops will not have the other variants
like `..end`, `start..` and `..`, which wouldn't have two separate
expressions, and the flag is simpler.
- Transforming an inclusive range to Hir and back makes it
non-inclusive, but I think that's fine because the block created by
`ForRange::into_for` to iterate over the array is not turned back
either.


## Documentation\*

Check one:
- [ ] No documentation needed.
- [x] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
  • Loading branch information
aakoshh authored Oct 3, 2024
1 parent 2a0d211 commit bd861f2
Show file tree
Hide file tree
Showing 20 changed files with 204 additions and 34 deletions.
62 changes: 57 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,57 @@ impl LValue {
}
}

#[derive(Debug, PartialEq, Eq, Clone)]
pub struct ForBounds {
pub start: Expression,
pub end: Expression,
pub inclusive: bool,
}

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) {
let end = if self.inclusive {
let end_span = self.end.span;
let end = ExpressionKind::Infix(Box::new(InfixExpression {
lhs: self.end,
operator: Spanned::from(end_span, BinaryOpKind::Add),
rhs: Expression::new(ExpressionKind::integer(FieldElement::from(1u32)), end_span),
}));
Expression::new(end, end_span)
} else {
self.end
};

(self.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, inclusive: false })
}

/// Create a range bounded inclusively below and above.
pub fn range_inclusive(start: Expression, end: Expression) -> Self {
Self::Range(ForBounds { start, end, inclusive: 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 +924,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 +1054,14 @@ 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(bounds) => {
format!(
"{}{}{}",
bounds.start,
if bounds.inclusive { "..=" } else { ".." },
bounds.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, inclusive: _ }) => {
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
14 changes: 9 additions & 5 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,13 @@ fn remove_interned_in_statement_kind(
}),
StatementKind::For(for_loop) => StatementKind::For(ForLoopStatement {
range: match for_loop.range {
ForRange::Range(from, to) => ForRange::Range(
remove_interned_in_expression(interner, from),
remove_interned_in_expression(interner, to),
),
ForRange::Range(ForBounds { start, end, inclusive }) => {
ForRange::Range(ForBounds {
start: remove_interned_in_expression(interner, start),
end: remove_interned_in_expression(interner, end),
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();
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() {
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
33 changes: 25 additions & 8 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 @@ -703,8 +719,8 @@ mod tests {
use crate::token::{CustomAttribute, FunctionAttribute, SecondaryAttribute, TestScope};

#[test]
fn test_single_double_char() {
let input = "! != + ( ) { } [ ] | , ; : :: < <= > >= & - -> . .. % / * = == << >>";
fn test_single_multi_char() {
let input = "! != + ( ) { } [ ] | , ; : :: < <= > >= & - -> . .. ..= % / * = == << >>";

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
Loading

0 comments on commit bd861f2

Please sign in to comment.