feat: while statement (#7280)

Co-authored-by: Michael J Klein <michaeljklein@users.noreply.github.com>
Co-authored-by: Akosh Farkash <aakoshh@gmail.com>
Co-authored-by: Tom French <15848336+TomAFrench@users.noreply.github.com>
Co-authored-by: Tom French <tom@tomfren.ch>

compiler/noirc_evaluator/src/ssa/ssa_gen/mod.rs

@@ -11,7 +11,7 @@ use iter_extended::{try_vecmap, vecmap};
 use noirc_errors::Location;
 use noirc_frontend::ast::{UnaryOp, Visibility};
 use noirc_frontend::hir_def::types::Type as HirType;
-use noirc_frontend::monomorphization::ast::{self, Expression, Program};
+use noirc_frontend::monomorphization::ast::{self, Expression, Program, While};
 
 use crate::{
     errors::RuntimeError,
@@ -153,6 +153,7 @@ impl<'a> FunctionContext<'a> {
             Expression::Cast(cast) => self.codegen_cast(cast),
             Expression::For(for_expr) => self.codegen_for(for_expr),
             Expression::Loop(block) => self.codegen_loop(block),
+            Expression::While(while_) => self.codegen_while(while_),
             Expression::If(if_expr) => self.codegen_if(if_expr),
             Expression::Tuple(tuple) => self.codegen_tuple(tuple),
             Expression::ExtractTupleField(tuple, index) => {
@@ -588,7 +589,7 @@ impl<'a> FunctionContext<'a> {
         Ok(Self::unit_value())
     }
 
-    /// Codegens a loop, creating three new blocks in the process.
+    /// Codegens a loop, creating two new blocks in the process.
     /// The return value of a loop is always a unit literal.
     ///
     /// For example, the loop `loop { body }` is codegen'd as:
@@ -620,6 +621,47 @@ impl<'a> FunctionContext<'a> {
         Ok(Self::unit_value())
     }
 
+    /// Codegens a while loop, creating three new blocks in the process.
+    /// The return value of a while is always a unit literal.
+    ///
+    /// For example, the loop `while cond { body }` is codegen'd as:
+    ///
+    /// ```text
+    ///   jmp while_entry()
+    /// while_entry:
+    ///   v0 = ... codegen cond ...
+    ///   jmpif v0, then: while_body, else: while_end  
+    /// while_body():
+    ///   v3 = ... codegen body ...
+    ///   jmp while_entry()
+    /// while_end():
+    ///   ... This is the current insert point after codegen_while finishes ...
+    /// ```
+    fn codegen_while(&mut self, while_: &While) -> Result<Values, RuntimeError> {
+        let while_entry = self.builder.insert_block();
+        let while_body = self.builder.insert_block();
+        let while_end = self.builder.insert_block();
+
+        self.builder.terminate_with_jmp(while_entry, vec![]);
+
+        // Codegen the entry (where the condition is)
+        self.builder.switch_to_block(while_entry);
+        let condition = self.codegen_non_tuple_expression(&while_.condition)?;
+        self.builder.terminate_with_jmpif(condition, while_body, while_end);
+
+        self.enter_loop(Loop { loop_entry: while_entry, loop_index: None, loop_end: while_end });
+
+        // Codegen the body
+        self.builder.switch_to_block(while_body);
+        self.codegen_expression(&while_.body)?;
+        self.builder.terminate_with_jmp(while_entry, vec![]);
+
+        // Finish by switching to the end of the while
+        self.builder.switch_to_block(while_end);
+        self.exit_loop();
+        Ok(Self::unit_value())
+    }
+
     /// Codegens an if expression, handling the case of what to do if there is no 'else'.
     ///
     /// For example, the expression `if cond { a } else { b }` is codegen'd as:

compiler/noirc_frontend/src/ast/statement.rs

@@ -46,6 +46,7 @@ pub enum StatementKind {
     Assign(AssignStatement),
     For(ForLoopStatement),
     Loop(Expression, Span /* loop keyword span */),
+    While(WhileStatement),
     Break,
     Continue,
     /// This statement should be executed at compile-time
@@ -103,11 +104,8 @@ impl StatementKind {
                     statement.add_semicolon(semi, span, last_statement_in_block, emit_error);
                 StatementKind::Comptime(statement)
             }
-            // A semicolon on a for loop is optional and does nothing
-            StatementKind::For(_) => self,
-
-            // A semicolon on a loop is optional and does nothing
-            StatementKind::Loop(..) => self,
+            // A semicolon on a for loop, loop or while is optional and does nothing
+            StatementKind::For(_) | StatementKind::Loop(..) | StatementKind::While(..) => self,
 
             // No semicolon needed for a resolved statement
             StatementKind::Interned(_) => self,
@@ -856,6 +854,13 @@ pub struct ForLoopStatement {
     pub span: Span,
 }
 
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct WhileStatement {
+    pub condition: Expression,
+    pub body: Expression,
+    pub while_keyword_span: Span,
+}
+
 impl Display for StatementKind {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
@@ -864,6 +869,9 @@ impl Display for StatementKind {
             StatementKind::Assign(assign) => assign.fmt(f),
             StatementKind::For(for_loop) => for_loop.fmt(f),
             StatementKind::Loop(block, _) => write!(f, "loop {}", block),
+            StatementKind::While(while_) => {
+                write!(f, "while {} {}", while_.condition, while_.body)
+            }
             StatementKind::Break => write!(f, "break"),
             StatementKind::Continue => write!(f, "continue"),
             StatementKind::Comptime(statement) => write!(f, "comptime {}", statement.kind),

compiler/noirc_frontend/src/ast/visitor.rs

@@ -310,6 +310,10 @@ pub trait Visitor {
         true
     }
 
+    fn visit_while_statement(&mut self, _condition: &Expression, _body: &Expression) -> bool {
+        true
+    }
+
     fn visit_comptime_statement(&mut self, _: &Statement) -> bool {
         true
     }
@@ -1165,6 +1169,12 @@ impl Statement {
                     block.accept(visitor);
                 }
             }
+            StatementKind::While(while_) => {
+                if visitor.visit_while_statement(&while_.condition, &while_.body) {
+                    while_.condition.accept(visitor);
+                    while_.body.accept(visitor);
+                }
+            }
             StatementKind::Comptime(statement) => {
                 if visitor.visit_comptime_statement(statement) {
                     statement.accept(visitor);

compiler/noirc_frontend/src/elaborator/lints.rs

@@ -283,6 +283,7 @@ fn can_return_without_recursing(interner: &NodeInterner, func_id: FuncId, expr_i
             // Rust doesn't seem to check the for loop body (it's bounds might mean it's never called).
             HirStatement::For(e) => check(e.start_range) && check(e.end_range),
             HirStatement::Loop(e) => check(e),
+            HirStatement::While(condition, block) => check(condition) && check(block),
             HirStatement::Comptime(_)
             | HirStatement::Break
             | HirStatement::Continue

compiler/noirc_frontend/src/elaborator/statements.rs

@@ -3,7 +3,7 @@ use noirc_errors::{Location, Span};
 use crate::{
     ast::{
         AssignStatement, Expression, ForLoopStatement, ForRange, Ident, ItemVisibility, LValue,
-        LetStatement, Path, Statement, StatementKind,
+        LetStatement, Path, Statement, StatementKind, WhileStatement,
     },
     hir::{
         resolution::{
@@ -37,6 +37,7 @@ impl<'context> Elaborator<'context> {
             StatementKind::Assign(assign) => self.elaborate_assign(assign),
             StatementKind::For(for_stmt) => self.elaborate_for(for_stmt),
             StatementKind::Loop(block, span) => self.elaborate_loop(block, span),
+            StatementKind::While(while_) => self.elaborate_while(while_),
             StatementKind::Break => self.elaborate_jump(true, statement.span),
             StatementKind::Continue => self.elaborate_jump(false, statement.span),
             StatementKind::Comptime(statement) => self.elaborate_comptime_statement(*statement),
@@ -258,6 +259,35 @@ impl<'context> Elaborator<'context> {
         (statement, Type::Unit)
     }
 
+    pub(super) fn elaborate_while(&mut self, while_: WhileStatement) -> (HirStatement, Type) {
+        let in_constrained_function = self.in_constrained_function();
+        if in_constrained_function {
+            self.push_err(ResolverError::WhileInConstrainedFn { span: while_.while_keyword_span });
+        }
+
+        let old_loop = std::mem::take(&mut self.current_loop);
+        self.current_loop = Some(Loop { is_for: false, has_break: false });
+        self.push_scope();
+
+        let condition_span = while_.condition.span;
+        let (condition, cond_type) = self.elaborate_expression(while_.condition);
+        let (block, _block_type) = self.elaborate_expression(while_.body);
+
+        self.unify(&cond_type, &Type::Bool, || TypeCheckError::TypeMismatch {
+            expected_typ: Type::Bool.to_string(),
+            expr_typ: cond_type.to_string(),
+            expr_span: condition_span,
+        });
+
+        self.pop_scope();
+
+        std::mem::replace(&mut self.current_loop, old_loop).expect("Expected a loop");
+
+        let statement = HirStatement::While(condition, block);
+
+        (statement, Type::Unit)
+    }
+
     fn elaborate_jump(&mut self, is_break: bool, span: noirc_errors::Span) -> (HirStatement, Type) {
         let in_constrained_function = self.in_constrained_function();
 

compiler/noirc_frontend/src/hir/comptime/display.rs

@@ -10,7 +10,7 @@ use crate::{
         ForBounds, ForLoopStatement, ForRange, GenericTypeArgs, IfExpression, IndexExpression,
         InfixExpression, LValue, Lambda, LetStatement, Literal, MatchExpression,
         MemberAccessExpression, MethodCallExpression, Pattern, PrefixExpression, Statement,
-        StatementKind, UnresolvedType, UnresolvedTypeData,
+        StatementKind, UnresolvedType, UnresolvedTypeData, WhileStatement,
     },
     hir_def::traits::TraitConstraint,
     node_interner::{InternedStatementKind, NodeInterner},
@@ -766,6 +766,11 @@ fn remove_interned_in_statement_kind(
         StatementKind::Loop(block, span) => {
             StatementKind::Loop(remove_interned_in_expression(interner, block), span)
         }
+        StatementKind::While(while_) => StatementKind::While(WhileStatement {
+            condition: remove_interned_in_expression(interner, while_.condition),
+            body: remove_interned_in_expression(interner, while_.body),
+            while_keyword_span: while_.while_keyword_span,
+        }),
         StatementKind::Comptime(statement) => {
             StatementKind::Comptime(Box::new(remove_interned_in_statement(interner, *statement)))
         }

compiler/noirc_frontend/src/hir/comptime/errors.rs

@@ -50,6 +50,10 @@ pub enum InterpreterError {
         typ: Type,
         location: Location,
     },
+    NonBoolUsedInWhile {
+        typ: Type,
+        location: Location,
+    },
     NonBoolUsedInConstrain {
         typ: Type,
         location: Location,
@@ -285,6 +289,7 @@ impl InterpreterError {
             | InterpreterError::ErrorNodeEncountered { location, .. }
             | InterpreterError::NonFunctionCalled { location, .. }
             | InterpreterError::NonBoolUsedInIf { location, .. }
+            | InterpreterError::NonBoolUsedInWhile { location, .. }
             | InterpreterError::NonBoolUsedInConstrain { location, .. }
             | InterpreterError::FailingConstraint { location, .. }
             | InterpreterError::NoMethodFound { location, .. }
@@ -413,6 +418,11 @@ impl<'a> From<&'a InterpreterError> for CustomDiagnostic {
                 let secondary = "If conditions must be a boolean value".to_string();
                 CustomDiagnostic::simple_error(msg, secondary, location.span)
             }
+            InterpreterError::NonBoolUsedInWhile { typ, location } => {
+                let msg = format!("Expected a `bool` but found `{typ}`");
+                let secondary = "While conditions must be a boolean value".to_string();
+                CustomDiagnostic::simple_error(msg, secondary, location.span)
+            }
             InterpreterError::NonBoolUsedInConstrain { typ, location } => {
                 let msg = format!("Expected a `bool` but found `{typ}`");
                 CustomDiagnostic::simple_error(msg, String::new(), location.span)

compiler/noirc_frontend/src/hir/comptime/hir_to_display_ast.rs

@@ -6,7 +6,7 @@ use crate::ast::{
     ConstructorExpression, ExpressionKind, ForLoopStatement, ForRange, GenericTypeArgs, Ident,
     IfExpression, IndexExpression, InfixExpression, LValue, Lambda, Literal,
     MemberAccessExpression, MethodCallExpression, Path, PathKind, PathSegment, Pattern,
-    PrefixExpression, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression,
+    PrefixExpression, UnresolvedType, UnresolvedTypeData, UnresolvedTypeExpression, WhileStatement,
 };
 use crate::ast::{ConstrainExpression, Expression, Statement, StatementKind};
 use crate::hir_def::expr::{
@@ -46,6 +46,11 @@ impl HirStatement {
                 span,
             }),
             HirStatement::Loop(block) => StatementKind::Loop(block.to_display_ast(interner), span),
+            HirStatement::While(condition, block) => StatementKind::While(WhileStatement {
+                condition: condition.to_display_ast(interner),
+                body: block.to_display_ast(interner),
+                while_keyword_span: span,
+            }),
             HirStatement::Break => StatementKind::Break,
             HirStatement::Continue => StatementKind::Continue,
             HirStatement::Expression(expr) => {

compiler/noirc_frontend/src/hir/comptime/interpreter.rs

@@ -1516,7 +1516,7 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
         let condition = match self.evaluate(if_.condition)? {
             Value::Bool(value) => value,
             value => {
-                let location = self.elaborator.interner.expr_location(&id);
+                let location = self.elaborator.interner.expr_location(&if_.condition);
                 let typ = value.get_type().into_owned();
                 return Err(InterpreterError::NonBoolUsedInIf { typ, location });
             }
@@ -1571,6 +1571,7 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
             HirStatement::Assign(assign) => self.evaluate_assign(assign),
             HirStatement::For(for_) => self.evaluate_for(for_),
             HirStatement::Loop(expression) => self.evaluate_loop(expression),
+            HirStatement::While(condition, block) => self.evaluate_while(condition, block),
             HirStatement::Break => self.evaluate_break(statement),
             HirStatement::Continue => self.evaluate_continue(statement),
             HirStatement::Expression(expression) => self.evaluate(expression),
@@ -1808,6 +1809,55 @@ impl<'local, 'interner> Interpreter<'local, 'interner> {
         result
     }
 
+    fn evaluate_while(&mut self, condition: ExprId, block: ExprId) -> IResult<Value> {
+        let was_in_loop = std::mem::replace(&mut self.in_loop, true);
+        let in_lsp = self.elaborator.interner.is_in_lsp_mode();
+        let mut counter = 0;
+        let mut result = Ok(Value::Unit);
+
+        loop {
+            let condition = match self.evaluate(condition)? {
+                Value::Bool(value) => value,
+                value => {
+                    let location = self.elaborator.interner.expr_location(&condition);
+                    let typ = value.get_type().into_owned();
+                    return Err(InterpreterError::NonBoolUsedInWhile { typ, location });
+                }
+            };
+            if !condition {
+                break;
+            }
+
+            self.push_scope();
+
+            let must_break = match self.evaluate(block) {
+                Ok(_) => false,
+                Err(InterpreterError::Break) => true,
+                Err(InterpreterError::Continue) => false,
+                Err(error) => {
+                    result = Err(error);
+                    true
+                }
+            };
+
+            self.pop_scope();
+
+            if must_break {
+                break;
+            }
+
+            counter += 1;
+            if in_lsp && counter == 10_000 {
+                let location = self.elaborator.interner.expr_location(&block);
+                result = Err(InterpreterError::LoopHaltedForUiResponsiveness { location });
+                break;
+            }
+        }
+
+        self.in_loop = was_in_loop;
+        result
+    }
+
     fn evaluate_break(&mut self, id: StmtId) -> IResult<Value> {
         if self.in_loop {
             Err(InterpreterError::Break)

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

@@ -102,6 +102,8 @@ pub enum ResolverError {
     LoopInConstrainedFn { span: Span },
     #[error("`loop` must have at least one `break` in it")]
     LoopWithoutBreak { span: Span },
+    #[error("`while` is only allowed in unconstrained functions")]
+    WhileInConstrainedFn { span: Span },
     #[error("break/continue are only allowed within loops")]
     JumpOutsideLoop { is_break: bool, span: Span },
     #[error("Only `comptime` globals can be mutable")]
@@ -442,7 +444,7 @@ impl<'a> From<&'a ResolverError> for Diagnostic {
             },
             ResolverError::LoopInConstrainedFn { span } => {
                 Diagnostic::simple_error(
-                    "loop is only allowed in unconstrained functions".into(),
+                    "`loop` is only allowed in unconstrained functions".into(),
                     "Constrained code must always have a known number of loop iterations".into(),
                     *span,
                 )
@@ -454,6 +456,13 @@ impl<'a> From<&'a ResolverError> for Diagnostic {
                     *span,
                 )
             },
+            ResolverError::WhileInConstrainedFn { span } => {
+                Diagnostic::simple_error(
+                    "`while` is only allowed in unconstrained functions".into(),
+                    "Constrained code must always have a known number of loop iterations".into(),
+                    *span,
+                )
+            },
             ResolverError::JumpOutsideLoop { is_break, span } => {
                 let item = if *is_break { "break" } else { "continue" };
                 Diagnostic::simple_error(

compiler/noirc_frontend/src/hir_def/stmt.rs

@@ -15,6 +15,7 @@ pub enum HirStatement {
     Assign(HirAssignStatement),
     For(HirForStatement),
     Loop(ExprId),
+    While(ExprId, ExprId),
     Break,
     Continue,
     Expression(ExprId),