feat: Improved error message for unexpected return type

crates/noirc_frontend/src/ast/expression.rs

@@ -364,11 +364,19 @@ pub struct FunctionDefinition {
     pub body: BlockExpression,
     pub span: Span,
     pub where_clause: Vec<TraitConstraint>,
-    pub return_type: UnresolvedType,
+    pub return_type: FunctionReturnType,
     pub return_visibility: noirc_abi::AbiVisibility,
     pub return_distinctness: noirc_abi::AbiDistinctness,
 }
 
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub enum FunctionReturnType {
+    /// Returns type is not specified.
+    Default(Span),
+    /// Everything else.
+    Ty(UnresolvedType, Span),
+}
+
 /// Describes the types of smart contract functions that are allowed.
 /// - All Noir programs in the non-contract context can be seen as `Secret`.
 #[derive(serde::Serialize, serde::Deserialize, Debug, Copy, Clone, PartialEq, Eq)]
@@ -636,3 +644,12 @@ impl Display for FunctionDefinition {
         )
     }
 }
+
+impl Display for FunctionReturnType {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            FunctionReturnType::Default(_) => f.write_str(""),
+            FunctionReturnType::Ty(ty, _) => write!(f, "{ty}"),
+        }
+    }
+}

crates/noirc_frontend/src/ast/function.rs

@@ -1,6 +1,6 @@
 use std::fmt::Display;
 
-use crate::{token::Attribute, Ident, Pattern};
+use crate::{token::Attribute, FunctionReturnType, Ident, Pattern};
 
 use super::{FunctionDefinition, UnresolvedType};
 
@@ -41,7 +41,10 @@ impl NoirFunction {
     }
 
     pub fn return_type(&self) -> UnresolvedType {
-        self.def.return_type.clone()
+        match &self.def.return_type {
+            FunctionReturnType::Default(_) => UnresolvedType::Unit,
+            FunctionReturnType::Ty(ty, _) => ty.clone(),
+        }
     }
     pub fn name(&self) -> &str {
         &self.name_ident().0.contents

crates/noirc_frontend/src/ast/traits.rs

@@ -3,7 +3,10 @@ use std::fmt::Display;
 use iter_extended::vecmap;
 use noirc_errors::Span;
 
-use crate::{BlockExpression, Expression, Ident, NoirFunction, UnresolvedGenerics, UnresolvedType};
+use crate::{
+    BlockExpression, Expression, FunctionReturnType, Ident, NoirFunction, UnresolvedGenerics,
+    UnresolvedType,
+};
 
 /// AST node for trait definitions:
 /// `trait name<generics> { ... items ... }`
@@ -24,7 +27,7 @@ pub enum TraitItem {
         name: Ident,
         generics: Vec<Ident>,
         parameters: Vec<(Ident, UnresolvedType)>,
-        return_type: UnresolvedType,
+        return_type: FunctionReturnType,
         where_clause: Vec<TraitConstraint>,
         body: Option<BlockExpression>,
     },

crates/noirc_frontend/src/hir/resolution/resolver.rs

@@ -726,6 +726,7 @@
             location,
             typ,
             parameters: parameters.into(),
+            return_type: func.def.return_type.clone(),
             return_visibility: func.def.return_visibility,
             return_distinctness: func.def.return_distinctness,
             has_body: !func.def.body.is_empty(),
@@ -1930,7 +1931,7 @@
                 let new_val = 10;
                 println(f"randomstring{new_val}{new_val}");
             }
             fn println<T>(x : T) -> T {
                 x
             }
         "#;

crates/noirc_frontend/src/hir/type_check/errors.rs

@@ -5,6 +5,7 @@ use thiserror::Error;
 use crate::hir::resolution::errors::ResolverError;
 use crate::hir_def::expr::HirBinaryOp;
 use crate::hir_def::types::Type;
+use crate::FunctionReturnType;
 use crate::Signedness;
 
 #[derive(Error, Debug, Clone, PartialEq, Eq)]
@@ -23,6 +24,8 @@ pub enum Source {
     Comparison,
     #[error("BinOp")]
     BinOp,
+    #[error("Return")]
+    Return(FunctionReturnType, Span),
 }
 
 #[derive(Error, Debug, Clone, PartialEq, Eq)]
@@ -195,6 +198,21 @@ impl From<TypeCheckError> for Diagnostic {
                     Source::StringLen => format!("Can only compare strings of the same length. Here LHS is of length {lhs}, and RHS is {rhs}"),
                     Source::Comparison => format!("Unsupported types for comparison: {lhs} and {rhs}"),
                     Source::BinOp => format!("Unsupported types for binary operation: {lhs} and {rhs}"),
+                    Source::Return(ret_ty, expr_span) => {
+                        let ret_ty_span = match ret_ty {
+                            FunctionReturnType::Default(span) | FunctionReturnType::Ty(_, span) => span
+                        };
+
+                        let mut diagnostic = Diagnostic::simple_error(format!("expected type {lhs}, found type {rhs}"), format!("expected {lhs} because of return type"), ret_ty_span);
+
+                        if let FunctionReturnType::Default(_) = ret_ty {
+                            diagnostic.add_note(format!("help: try adding a return type: `-> {rhs}`"));
+                        }
+
+                        diagnostic.add_secondary(format!("{rhs} returned here"), expr_span);
+
+                        return diagnostic
+                    },
                 };
 
                 Diagnostic::simple_error(message, String::new(), span)

crates/noirc_frontend/src/hir/type_check/mod.rs

@@ -14,10 +14,13 @@ mod stmt;
 pub use errors::TypeCheckError;
 
 use crate::{
+    hir_def::{expr::HirExpression, stmt::HirStatement},
     node_interner::{ExprId, FuncId, NodeInterner, StmtId},
     Type,
 };
 
+use self::errors::Source;
+
 type TypeCheckFn = Box<dyn FnOnce() -> Result<(), TypeCheckError>>;
 
 pub struct TypeChecker<'interner> {
@@ -57,23 +60,58 @@ pub fn type_check_func(interner: &mut NodeInterner, func_id: FuncId) -> Vec<Type
 
     // Check declared return type and actual return type
     if !can_ignore_ret {
+        let (expr_span, empty_function) = function_info(interner, function_body_id);
+
         let func_span = interner.expr_span(function_body_id); // XXX: We could be more specific and return the span of the last stmt, however stmts do not have spans yet
         function_last_type.unify_with_coercions(
             &declared_return_type,
             *function_body_id,
             interner,
             &mut errors,
-            || TypeCheckError::TypeMismatch {
-                expected_typ: declared_return_type.to_string(),
-                expr_typ: function_last_type.to_string(),
-                expr_span: func_span,
+            || {
+                let mut error = TypeCheckError::TypeMismatchWithSource {
+                    lhs: declared_return_type.clone(),
+                    rhs: function_last_type.clone(),
+                    span: func_span,
+                    source: Source::Return(meta.return_type, expr_span),
+                };
+
+                if empty_function {
+                    error = error.add_context(
+                        "implicitly returns `()` as its body has no tail or `return` expression",
+                    );
+                }
+
+                error
             },
         );
     }
 
     errors
 }
 
+fn function_info(
+    interner: &mut NodeInterner,
+    function_body_id: &ExprId,
+) -> (noirc_errors::Span, bool) {
+    let (expr_span, empty_function) =
+        if let HirExpression::Block(block) = interner.expression(function_body_id) {
+            let last_stmt = block.statements().last();
+            let mut span = interner.expr_span(function_body_id);
+
+            if let Some(last_stmt) = last_stmt {
+                if let HirStatement::Expression(expr) = interner.statement(last_stmt) {
+                    span = interner.expr_span(&expr);
+                }
+            }
+
+            (span, last_stmt.is_none())
+        } else {
+            (interner.expr_span(function_body_id), false)
+        };
+    (expr_span, empty_function)
+}
+
 impl<'interner> TypeChecker<'interner> {
     fn new(interner: &'interner mut NodeInterner) -> Self {
         Self { delayed_type_checks: Vec::new(), interner, errors: vec![] }
@@ -149,14 +187,14 @@ mod test {
         stmt::HirStatement,
     };
     use crate::node_interner::{DefinitionKind, FuncId, NodeInterner};
-    use crate::BinaryOpKind;
     use crate::{
         hir::{
             def_map::{CrateDefMap, LocalModuleId, ModuleDefId},
             resolution::{path_resolver::PathResolver, resolver::Resolver},
         },
         parse_program, FunctionKind, Path,
     };
+    use crate::{BinaryOpKind, FunctionReturnType};
 
     #[test]
     fn basic_let() {
@@ -237,6 +275,7 @@ mod test {
             return_visibility: noirc_abi::AbiVisibility::Private,
             return_distinctness: noirc_abi::AbiDistinctness::DuplicationAllowed,
             has_body: true,
+            return_type: FunctionReturnType::Default(Span::default()),
         };
         interner.push_fn_meta(func_meta, func_id);
 

crates/noirc_frontend/src/hir_def/function.rs

@@ -7,7 +7,7 @@ use super::stmt::HirPattern;
 use crate::hir::def_map::ModuleId;
 use crate::node_interner::{ExprId, NodeInterner};
 use crate::{token::Attribute, FunctionKind};
-use crate::{ContractFunctionType, Type};
+use crate::{ContractFunctionType, FunctionReturnType, Type};
 
 /// A Hir function is a block expression
 /// with a list of statements
@@ -137,6 +137,8 @@ pub struct FuncMeta {
 
     pub parameters: Parameters,
 
+    pub return_type: FunctionReturnType,
+
     pub return_visibility: AbiVisibility,
 
     pub return_distinctness: AbiDistinctness,

crates/noirc_frontend/src/parser/parser.rs

@@ -23,6 +23,7 @@
 //! prevent other parsers from being tried afterward since there is no longer an error. Thus, they should
 //! be limited to cases like the above `fn` example where it is clear we shouldn't back out of the
 //! current parser to try alternative parsers in a `choice` expression.
+use super::spanned;
 use super::{
     foldl_with_span, labels::ParsingRuleLabel, parameter_name_recovery, parameter_recovery,
     parenthesized, then_commit, then_commit_ignore, top_level_statement_recovery, ExprParser,
@@ -34,10 +35,11 @@
 use crate::parser::{force, ignore_then_commit, statement_recovery};
 use crate::token::{Attribute, Keyword, Token, TokenKind};
 use crate::{
-    BinaryOp, BinaryOpKind, BlockExpression, ConstrainStatement, FunctionDefinition, Ident,
-    IfExpression, InfixExpression, LValue, Lambda, Literal, NoirFunction, NoirStruct, NoirTrait,
-    NoirTypeAlias, Path, PathKind, Pattern, Recoverable, TraitConstraint, TraitImpl, TraitImplItem,
-    TraitItem, TypeImpl, UnaryOp, UnresolvedTypeExpression, UseTree, UseTreeKind,
+    BinaryOp, BinaryOpKind, BlockExpression, ConstrainStatement, FunctionDefinition,
+    FunctionReturnType, Ident, IfExpression, InfixExpression, LValue, Lambda, Literal,
+    NoirFunction, NoirStruct, NoirTrait, NoirTypeAlias, Path, PathKind, Pattern, Recoverable,
+    TraitConstraint, TraitImpl, TraitImplItem, TraitItem, TypeImpl, UnaryOp,
+    UnresolvedTypeExpression, UseTree, UseTreeKind,
 };
 
 use chumsky::prelude::*;
@@ -258,17 +260,19 @@
         .map(|ret| ret.unwrap_or(UnresolvedType::Unspecified))
 }
 
-fn function_return_type() -> impl NoirParser<((AbiDistinctness, AbiVisibility), UnresolvedType)> {
+fn function_return_type() -> impl NoirParser<((AbiDistinctness, AbiVisibility), FunctionReturnType)>
+{
     just(Token::Arrow)
         .ignore_then(optional_distinctness())
         .then(optional_visibility())
-        .then(parse_type())
+        .then(spanned(parse_type()))
         .or_not()
-        .map(|ret| {
-            ret.unwrap_or((
+        .map_with_span(|ret, span| match ret {
+            Some((head, (ty, span))) => (head, FunctionReturnType::Ty(ty, span)),
+            None => (
                 (AbiDistinctness::DuplicationAllowed, AbiVisibility::Private),
-                UnresolvedType::Unit,
-            ))
+                FunctionReturnType::Default(span),
+            ),
         })
 }
 
@@ -329,10 +333,10 @@
 }
 
 fn self_parameter() -> impl NoirParser<(Pattern, UnresolvedType, AbiVisibility)> {
     let refmut_pattern = just(Token::Ampersand).then_ignore(keyword(Keyword::Mut));
     let mut_pattern = keyword(Keyword::Mut);
 
     refmut_pattern
         .or(mut_pattern)
         .map_with_span(|token, span| (token, span))
         .or_not()
@@ -773,7 +777,7 @@
     let shorthand_operators = right_shift_operator().or(one_of(shorthand_operators));
     let shorthand_syntax = shorthand_operators.then_ignore(just(Token::Assign));
 
     // Since >> is lexed as two separate greater-thans, >>= is lexed as > >=, so
     // we need to account for that case here as well.
     let right_shift_fix =
         just(Token::Greater).then(just(Token::GreaterEqual)).map(|_| Token::ShiftRight);
@@ -801,7 +805,7 @@
 
     let dereferences = just(Token::Star).repeated();
 
     let lvalues =
         l_ident.then(l_member_rhs.or(l_index).repeated()).foldl(|lvalue, rhs| match rhs {
             LValueRhs::MemberAccess(field_name) => {
                 LValue::MemberAccess { object: Box::new(lvalue), field_name }
@@ -809,7 +813,7 @@
             LValueRhs::Index(index) => LValue::Index { array: Box::new(lvalue), index },
         });
 
     dereferences.then(lvalues).foldr(|_, lvalue| LValue::Dereference(Box::new(lvalue)))
 }
 
 fn parse_type<'a>() -> impl NoirParser<UnresolvedType> + 'a {