Constrain early returns in functions in addition to closures

roc-lang · Nov 8, 2024 · 9a2b99c · 9a2b99c
1 parent dfb2966
commit 9a2b99c
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 0 deletions.
diff --git a/crates/compiler/constrain/src/expr.rs b/crates/compiler/constrain/src/expr.rs
@@ -2114,6 +2114,21 @@ fn constrain_function_def(
                 ret_type_index,
             ));
 
+            let mut early_return_constraints = Vec::with_capacity(function_def.early_returns.len());
+            for (early_return_variable, early_return_region) in &function_def.early_returns {
+                let early_return_var = constraints.push_variable(*early_return_variable);
+                let early_return_con = constraints.equal_types(
+                    early_return_var,
+                    return_type_annotation_expected,
+                    Category::Return,
+                    *early_return_region,
+                );
+
+                early_return_constraints.push(early_return_con);
+            }
+
+            let early_returns_constraint = constraints.and_constraint(early_return_constraints);
+
             let solved_fn_type = {
                 // TODO(types-soa) optimize for Variable
                 let pattern_types = types.from_old_type_slice(
@@ -2151,6 +2166,7 @@ fn constrain_function_def(
                     std::file!(),
                     std::line!(),
                 ),
+                early_returns_constraint,
                 constraints.let_constraint(
                     [],
                     argument_pattern_state.vars,

diff --git a/crates/compiler/test_gen/src/gen_return.rs b/crates/compiler/test_gen/src/gen_return.rs
@@ -107,3 +107,45 @@ fn early_return_solo() {
         true
     );
 }
+
+#[test]
+#[cfg(any(feature = "gen-llvm", feature = "gen-wasm", feature = "gen-dev"))]
+fn early_return_annotated_function() {
+    assert_evals_to!(
+        indoc!(
+            r#"
+            app "test" provides [main] to "./platform"
+
+            failIfLessThanFive : U64 -> Result {} [LessThanFive]
+            failIfLessThanFive = \n ->
+                if n < 5 then
+                    Err LessThanFive
+                else
+                    Ok {}
+
+            validateInput : Str -> Result U64 [InvalidNumStr, LessThanFive]
+            validateInput = \str ->
+                num = try Str.toU64 str
+
+                when failIfLessThanFive num is
+                    Err err ->
+                        return Err err
+
+                    Ok {} ->
+                        Ok num
+
+            main : List Str
+            main =
+                ["abc", "3", "7"]
+                |> List.map validateInput
+                |> List.map Inspect.toStr
+            "#
+        ),
+        RocList::from_slice(&[
+            RocStr::from("(Err InvalidNumStr)"),
+            RocStr::from("(Err LessThanFive)"),
+            RocStr::from("(Ok 7)")
+        ]),
+        RocList<RocStr>
+    );
+}