Check return expressions against the function's return type

src/checker.rs

@@ -55,6 +55,10 @@ pub struct Checker {
 
     /// Nesting depth of loops (>0 means we're inside a loop).
     loop_depth: usize,
+
+    /// Return types of the enclosing functions and lambdas, innermost last.
+    /// A `return` expression is constrained against the last entry.
+    ret_types: Vec<TypeID>,
 }
 
 /// Returns true if the type is or contains a borrowed type (`[T]` or `&T`).
@@ -187,6 +191,7 @@ impl Checker {
             constraints: vec![],
             errors: vec![],
             loop_depth: 0,
+            ret_types: vec![],
         }
     }
 
@@ -716,7 +721,31 @@ impl Checker {
                 ty
             }
             Expr::Arena(block) => self.check_expr(*block, arena, decls),
-            Expr::Return(expr) => self.check_expr(*expr, arena, decls),
+            Expr::Return(expr) => {
+                let ty = self.check_expr(*expr, arena, decls);
+
+                // Constrain against the enclosing function's return type.
+                // Without this, only a return in tail position is checked
+                // (via the body's type), so an early return with the wrong
+                // type reaches codegen and trips the backend's verifier.
+                match self.ret_types.last() {
+                    Some(&ret) => {
+                        self.eq(
+                            ty,
+                            ret,
+                            arena.locs[id],
+                            "return type must match function return type",
+                        );
+
+                        // Take on the enclosing return type rather than the
+                        // operand's, so a return in tail position doesn't
+                        // report the same mismatch twice (once here, once
+                        // from the body-level check below).
+                        ret
+                    }
+                    None => ty,
+                }
+            }
             Expr::Assume(cond) => {
                 self.check_expr(*cond, arena, decls);
                 mk_type(Type::Void)
@@ -955,13 +984,25 @@ impl Checker {
                     param_types.push(ty);
                 }
 
+                // The lambda's return type isn't known up front, so bind a
+                // fresh variable that both the body and any `return` inside
+                // it unify with.
+                let lambda_ret = self.fresh();
+                self.ret_types.push(lambda_ret);
                 let rt = self.check_expr(*body, arena, decls);
+                self.ret_types.pop();
+                self.eq(
+                    rt,
+                    lambda_ret,
+                    arena.locs[*body],
+                    "return type must match function return type",
+                );
 
                 while self.vars.len() > n {
                     self.vars.pop();
                 }
 
-                func(tuple(param_types), rt)
+                func(tuple(param_types), lambda_ret)
             }
             Expr::Error => self.fresh(),
         };
@@ -1081,7 +1122,9 @@ impl Checker {
             }
 
             // Check the body of the function.
+            self.ret_types.push(func_decl.ret);
             let ty = self.check_expr(body, &func_decl.arena, decls);
+            self.ret_types.pop();
 
             if func_decl.ret != mk_type(Type::Void) {
                 self.eq(

src/parser.rs

@@ -785,9 +785,10 @@ fn parse_stmt(arena: &mut ExprArena, typevars: &[Name], cx: &mut ParseContext) -
             arena.add(Expr::While(cond, body), cx.lex.loc)
         }
         Token::Return => {
+            let loc = cx.lex.loc;
             cx.next();
             let e = parse_expr(arena, typevars, cx);
-            arena.add(Expr::Return(e), cx.lex.loc)
+            arena.add(Expr::Return(e), loc)
         }
         Token::Break => {
             let loc = cx.lex.loc;

tests/cases/checker/early_return_ok.lyte

@@ -0,0 +1,21 @@
+// Well-typed early returns, in a value-returning function and in a lambda.
+// expected stdout:
+// compilation successful
+// 0
+// 1
+// 5
+// 3
+
+clamp01(x: i32) -> i32 {
+    if x < 0 { return 0 }
+    if x > 1 { return 1 }
+    x
+}
+
+main() {
+    print(clamp01(-4))
+    print(clamp01(7))
+    print(clamp01(5) + 4)
+    let f = |x: i32| { return x + 1 }
+    print(f(2))
+}

tests/cases/checker/early_return_type_mismatch.lyte

@@ -0,0 +1,16 @@
+// An early return nested in a branch must match the function's return type.
+// Only tail-position returns used to be checked, so this reached codegen and
+// tripped the Cranelift verifier.
+
+// args: --check
+// expected stdout:
+// ❌ ../tests/cases/checker/early_return_type_mismatch.lyte:12:12: return type must match function return type: f32 vs i32
+//         if b { return 1.5 }
+//                ^
+
+f(b: bool) -> i32 {
+    if b { return 1.5 }
+    0
+}
+
+main() { print(f(true)) }

tests/cases/checker/return_type_mismatch.lyte

@@ -1,8 +1,9 @@
 // args: --check
 // expected stdout:
-// ❌ ../tests/cases/checker/return_type_mismatch.lyte:6:18: return type must match function return type: f32 vs i32
-//     f(x: i32) -> i32 {
-//                      ^
+// ❌ ../tests/cases/checker/return_type_mismatch.lyte:8:5: return type must match function return type: f32 vs i32
+//         return 1.0
+//         ^
+
 f(x: i32) -> i32 {
     return 1.0
 }

tests/cases/checker/return_value_from_void.lyte

@@ -0,0 +1,14 @@
+// Returning a value from a function with no declared return type is an
+// error. It used to type-check and then trip the Cranelift verifier.
+
+// args: --check
+// expected stdout:
+// ❌ ../tests/cases/checker/return_value_from_void.lyte:11:5: return type must match function return type: i32 vs void
+//         return 0
+//         ^
+
+f() {
+    return 0
+}
+
+main() { f() }

tests/cases/generics/generic_global_return_infers.lyte

@@ -0,0 +1,25 @@
+// A `return` constrains its operand against the enclosing function's return
+// type, which is enough to infer T for a generic global used without ⟨T⟩ —
+// including for an early return inside a branch.
+// expected stdout:
+// compilation successful
+// 42
+// assert(true)
+// 0
+
+var pool<T>: [T; 4]
+
+get<T>(idx: i32) -> T {
+    if idx >= 0 && idx < 4 {
+        return pool[idx]
+    }
+    return pool[0]
+}
+
+main() {
+    pool⟨i32⟩[2] = 42
+    pool⟨f32⟩[2] = 1.5
+    print(get⟨i32⟩(2))
+    assert(get⟨f32⟩(2) == 1.5)
+    print(get⟨i32⟩(99))
+}

tests/cases/generics/unsolved_typevar_error.lyte

@@ -1,19 +1,20 @@
 // Using a generic global without ⟨T⟩ inside a generic function where
-// the return type alone can't determine T should be an error, not a hang.
+// nothing constrains T should be an error, not a hang.
 
 // args: --check
 // expected stdout:
-// ❌ ../tests/cases/generics/unsolved_typevar_error.lyte:14:16: ambiguous constraint: ?4 == ?4
-//             return pool[idx]
-//                    ^
+// ❌ ../tests/cases/generics/unsolved_typevar_error.lyte:16:13: ambiguous constraint: ?1 == ?1
+//         let x = pool[0]
+//                 ^
+// ❌ ../tests/cases/generics/unsolved_typevar_error.lyte:16:5: ambiguous constraint: ?1 == ?1
+//         let x = pool[0]
+//         ^
 
 var pool<T>: [T; 4]
 
-get<T>(idx: i32) -> T {
-    if idx >= 0 && idx < 4 {
-        return pool[idx]
-    }
-    return pool[0]
+count<T>() -> i32 {
+    let x = pool[0]
+    0
 }
 
 main {}