fix: require loop/for/while body to be unit

compiler/noirc_frontend/src/ast/expression.rs

@@ -245,6 +245,45 @@ impl Expression {
     pub fn new(kind: ExpressionKind, span: Span) -> Expression {
         Expression { kind, span }
     }
+
+    /// Returns the innermost span that gives this expression its type.
+    pub fn type_span(&self) -> Span {
+        match &self.kind {
+            ExpressionKind::Block(block_expression)
+            | ExpressionKind::Comptime(block_expression, _)
+            | ExpressionKind::Unsafe(block_expression, _) => {
+                if let Some(statement) = block_expression.statements.last() {
+                    statement.type_span()
+                } else {
+                    self.span
+                }
+            }
+            ExpressionKind::Parenthesized(expression) => expression.type_span(),
+            ExpressionKind::Literal(..)
+            | ExpressionKind::Prefix(..)
+            | ExpressionKind::Index(..)
+            | ExpressionKind::Call(..)
+            | ExpressionKind::MethodCall(..)
+            | ExpressionKind::Constrain(..)
+            | ExpressionKind::Constructor(..)
+            | ExpressionKind::MemberAccess(..)
+            | ExpressionKind::Cast(..)
+            | ExpressionKind::Infix(..)
+            | ExpressionKind::If(..)
+            | ExpressionKind::Match(..)
+            | ExpressionKind::Variable(..)
+            | ExpressionKind::Tuple(..)
+            | ExpressionKind::Lambda(..)
+            | ExpressionKind::Quote(..)
+            | ExpressionKind::Unquote(..)
+            | ExpressionKind::AsTraitPath(..)
+            | ExpressionKind::TypePath(..)
+            | ExpressionKind::Resolved(..)
+            | ExpressionKind::Interned(..)
+            | ExpressionKind::InternedStatement(..)
+            | ExpressionKind::Error => self.span,
+        }
+    }
 }
 
 pub type BinaryOp = Spanned<BinaryOpKind>;

compiler/noirc_frontend/src/ast/statement.rs

@@ -73,6 +73,24 @@ impl Statement {
         self.kind = self.kind.add_semicolon(semi, span, last_statement_in_block, emit_error);
         self
     }
+
+    /// Returns the innermost span that gives this statement its type.
+    pub fn type_span(&self) -> Span {
+        match &self.kind {
+            StatementKind::Expression(expression) => expression.type_span(),
+            StatementKind::Comptime(statement) => statement.type_span(),
+            StatementKind::Let(..)
+            | StatementKind::Assign(..)
+            | StatementKind::For(..)
+            | StatementKind::Loop(..)
+            | StatementKind::While(..)
+            | StatementKind::Break
+            | StatementKind::Continue
+            | StatementKind::Semi(..)
+            | StatementKind::Interned(..)
+            | StatementKind::Error => self.span,
+        }
+    }
 }
 
 impl StatementKind {

compiler/noirc_frontend/src/elaborator/enums.rs

@@ -274,7 +274,7 @@
             let columns = vec![Column::new(variable_to_match, pattern)];
 
             let guard = None;
-            let body_span = branch.span;
+            let body_span = branch.type_span();
             let (body, body_type) = self.elaborate_expression(branch);
 
             self.unify(&body_type, &result_type, || TypeCheckError::TypeMismatch {
@@ -291,7 +291,7 @@
 
     /// Convert an expression into a Pattern, defining any variables within.
     fn expression_to_pattern(&mut self, expression: Expression, expected_type: &Type) -> Pattern {
-        let expr_span = expression.span;
+        let expr_span = expression.type_span();
         let unify_with_expected_type = |this: &mut Self, actual| {
             this.unify(actual, expected_type, || TypeCheckError::TypeMismatch {
                 expected_typ: expected_type.to_string(),
@@ -647,7 +647,7 @@
 
     /// Compiles the cases and fallback cases for integer and range patterns.
     ///
     /// Integers have an infinite number of constructors, so we specialise the
     /// compilation of integer and range patterns.
     fn compile_int_cases(
         &mut self,

compiler/noirc_frontend/src/elaborator/expressions.rs

@@ -971,8 +971,8 @@ impl<'context> Elaborator<'context> {
         if_expr: IfExpression,
         target_type: Option<&Type>,
     ) -> (HirExpression, Type) {
-        let expr_span = if_expr.condition.span;
-        let consequence_span = if_expr.consequence.span;
+        let expr_span = if_expr.condition.type_span();
+        let consequence_span = if_expr.consequence.type_span();
         let (condition, cond_type) = self.elaborate_expression(if_expr.condition);
         let (consequence, mut ret_type) =
             self.elaborate_expression_with_target_type(if_expr.consequence, target_type);
@@ -984,9 +984,10 @@ impl<'context> Elaborator<'context> {
         });
 
         let (alternative, else_type, error_span) = if let Some(alternative) = if_expr.alternative {
+            let alternative_span = alternative.type_span();
             let (else_, else_type) =
                 self.elaborate_expression_with_target_type(alternative, target_type);
-            (Some(else_), else_type, expr_span)
+            (Some(else_), else_type, alternative_span)
         } else {
             (None, Type::Unit, consequence_span)
         };

compiler/noirc_frontend/src/elaborator/statements.rs

@@ -219,7 +219,14 @@ impl<'context> Elaborator<'context> {
 
         self.interner.push_definition_type(identifier.id, start_range_type);
 
-        let (block, _block_type) = self.elaborate_expression(block);
+        let block_span = block.type_span();
+        let (block, block_type) = self.elaborate_expression(block);
+
+        self.unify(&block_type, &Type::Unit, || TypeCheckError::TypeMismatch {
+            expected_typ: Type::Unit.to_string(),
+            expr_typ: block_type.to_string(),
+            expr_span: block_span,
+        });
 
         self.pop_scope();
         self.current_loop = old_loop;
@@ -244,7 +251,14 @@ impl<'context> Elaborator<'context> {
         self.current_loop = Some(Loop { is_for: false, has_break: false });
         self.push_scope();
 
-        let (block, _block_type) = self.elaborate_expression(block);
+        let block_span = block.type_span();
+        let (block, block_type) = self.elaborate_expression(block);
+
+        self.unify(&block_type, &Type::Unit, || TypeCheckError::TypeMismatch {
+            expected_typ: Type::Unit.to_string(),
+            expr_typ: block_type.to_string(),
+            expr_span: block_span,
+        });
 
         self.pop_scope();
 
@@ -269,16 +283,24 @@ impl<'context> Elaborator<'context> {
         self.current_loop = Some(Loop { is_for: false, has_break: false });
         self.push_scope();
 
-        let condition_span = while_.condition.span;
+        let condition_span = while_.condition.type_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,
         });
 
+        let block_span = while_.body.type_span();
+        let (block, block_type) = self.elaborate_expression(while_.body);
+
+        self.unify(&block_type, &Type::Unit, || TypeCheckError::TypeMismatch {
+            expected_typ: Type::Unit.to_string(),
+            expr_typ: block_type.to_string(),
+            expr_span: block_span,
+        });
+
         self.pop_scope();
 
         std::mem::replace(&mut self.current_loop, old_loop).expect("Expected a loop");

compiler/noirc_frontend/src/tests.rs

@@ -3190,7 +3190,7 @@
 }
 
 #[test]
 fn dont_infer_globals_to_u32_from_type_use() {
     let src = r#"
         global ARRAY_LEN = 3;
         global STR_LEN: _ = 2;
@@ -3220,7 +3220,7 @@
 }
 
 #[test]
 fn dont_infer_partial_global_types() {
     let src = r#"
         pub global ARRAY: [Field; _] = [0; 3];
         pub global NESTED_ARRAY: [[Field; _]; 3] = [[]; 3];
@@ -3463,7 +3463,7 @@
     // wouldn't panic due to infinite recursion, but the errors asserted here
     // come from the compilation checks, which does static analysis to catch the
     // problem before it even has a chance to cause a panic.
     let srcs = vec![
         r#"
         fn main() {
             main()
@@ -3525,7 +3525,7 @@
         "#,
     ];
 
     for src in srcs {
         let errors = get_program_errors(src);
         assert!(
             !errors.is_empty(),
@@ -3544,7 +3544,7 @@
 
 #[test]
 fn unconditional_recursion_pass() {
     let srcs = vec![
         r#"
         fn main() {
             if false { main(); }
@@ -3586,7 +3586,7 @@
         "#,
     ];
 
     for src in srcs {
         assert_no_errors(src);
     }
 }
@@ -4371,3 +4371,56 @@
     let src = "//\n".repeat(10_000);
     assert_no_errors(&src);
 }
+
+#[test]
+fn errors_if_for_body_type_is_not_unit() {
+    let src = r#"
+    fn main() {
+        for _ in 0..1 {
+            1
+        }
+    }
+    "#;
+    let errors = get_program_errors(src);
+    assert_eq!(errors.len(), 1);
+
+    let CompilationError::TypeError(TypeCheckError::TypeMismatch { .. }) = &errors[0].0 else {
+        panic!("Expected a TypeMismatch error");
+    };
+}
+
+#[test]
+fn errors_if_loop_body_type_is_not_unit() {
+    let src = r#"
+    unconstrained fn main() {
+        loop {
+            if false { break; }
+
+            1
+        }
+    }
+    "#;
+    let errors = get_program_errors(src);
+    assert_eq!(errors.len(), 1);
+
+    let CompilationError::TypeError(TypeCheckError::TypeMismatch { .. }) = &errors[0].0 else {
+        panic!("Expected a TypeMismatch error");
+    };
+}
+
+#[test]
+fn errors_if_while_body_type_is_not_unit() {
+    let src = r#"
+    unconstrained fn main() {
+        while 1 == 1 {
+            1
+        }
+    }
+    "#;
+    let errors = get_program_errors(src);
+    assert_eq!(errors.len(), 1);
+
+    let CompilationError::TypeError(TypeCheckError::TypeMismatch { .. }) = &errors[0].0 else {
+        panic!("Expected a TypeMismatch error");
+    };
+}