Constrain early returns in functions in addition to closures

[smores56]

When I initially implemented the return keyword, I only added type constraints to early returns in untyped closures because I misunderstood how typed closures/functions were constrained. This copies the code for early return constraints to the function type constraining code.

Note: if you annotate a function like myFunction : Result U64 [MyErr] -> Result U64 [MyErr, Other], the error types will not properly unify, and you'll get a type error. This is a behavior that exists without early returns (I checked), and you can find discussion from Ayaz surrounding it here.

[smores56]

I've been investigating this issue on and off in the past week, and I've found a small reproduction of the issue, which seems to be somewhere at or before monomorphization. With the following function:

validateInput : Str -> Result U64 _
validateInput = \str ->
    num =
        when Str.toU64 str is
            Ok a -> a
            Err err ->
                return Err err

    Ok num

validateInput "123"

If I comment out the type annotation, it compiles correctly. Of note is that the monomorphized early return branch of the when statement is

Let(`#UserApp.15`, Tag { tag_layout: NonRecursive([[InLayout(22)], [InLayout(U64)]]), tag_id: 0, arguments: [`#UserApp.n`], reuse: None }, InLayout(26), Ret(`#UserApp.15`))

Which defines #UserApp.15 = Err n and then returns #UserApp.15. If I keep the type annotation, we fail to properly monomorphize with the following error:

Full source: procedure `#UserApp.validateInput` (`#UserApp.str`):
    joinpoint `#UserApp.12` `#UserApp.num`:
        ret `#UserApp.num`;
    in
    let `#UserApp.11` : [C {}, C U64] = CallByName `Str.toU64` `#UserApp.str`;
    let `#UserApp.16` : U8 = 1i64;
    let `#UserApp.17` : U8 = GetTagId `#UserApp.11`;
    let `#UserApp.18` : Int1 = lowlevel Eq `#UserApp.16` `#UserApp.17`;
    if `#UserApp.18` then
        let `#UserApp.a` : U64 = UnionAtIndex (Id 1) (Index 0) `#UserApp.11`;
        jump `#UserApp.12` `#UserApp.a`;
    else
        let `#UserApp.err` : {} = UnionAtIndex (Id 0) (Index 0) `#UserApp.11`;
        ret `#UserApp.err`;
thread 'return_annotated' panicked at crates/compiler/test_mono/src/tests.rs:227:5:
IR problems found:
── SYMBOL LAYOUT DOESN'T MATCH ITS USE ─────────────────────────────────────────

in validateInput : (Str) -> U64 ((niche {}))

10│          jump `#UserApp.12` `#UserApp.a`;
11│>     else
12│          let `#UserApp.err` : {} = UnionAtIndex (Id 0) (Index 0) `#UserApp.11`;

#UserApp.err defined here with layout {}

11│      else
12│>         let `#UserApp.err` : {} = UnionAtIndex (Id 0) (Index 0) `#UserApp.11`;
13│          ret `#UserApp.err`;

#UserApp.err used as a return value here with layout U64

The monomorphized early return branch in the annotated version is just Ret(#UserApp.err), which is simpler. If we look at the full validateInput mono code for the non-annotated version, the difference is more clear:

procedure Test.1 (Test.2):
    joinpoint Test.12 Test.3:
        let Test.10 : [C {}, C U64] = TagId(1) Test.3;
        ret Test.10;
    in
    let Test.11 : [C {}, C U64] = CallByName Str.26 Test.2;
    let Test.16 : U8 = 1i64;
    let Test.17 : U8 = GetTagId Test.11;
    let Test.18 : Int1 = lowlevel Eq Test.16 Test.17;
    if Test.18 then
        let Test.4 : U64 = UnionAtIndex (Id 1) (Index 0) Test.11;
        jump Test.12 Test.4;
    else
        let Test.5 : {} = UnionAtIndex (Id 0) (Index 0) Test.11;
        let Test.15 : [C {}, C U64] = TagId(0) Test.5;
        ret Test.15;

Without a type annotation, we extract the payloads for each tag union from the when condition, and then wrap them in their respective tags. With a type annotation, we only extract the tag payloads, but we don't wrap into their tags.

I'll keep investigating this, my suspicion is that we are doing something different in type-constraining untyped functions vs. typed functions and that leads us to handle pattern binding extraction differently between these two cases.

[smores56]

Okay, it seems that in the above example, the return type is being inferred as U64 for the annotated example regardless of the annotated return type. The only value that has a definite type of U64 in this function body is a/num, since they are the successful result of Str.toU64. Presumably, something is going wrong in the type inference after getting the type of num and up to finishing inferring the closure's type.

[smores56]

Even more specifically, the type of values in early return statements are being totally ignored in type-annotated function definitions. The reason I thought that the return type was being inferred as U64 is because the type was being inferred as Result U64 [], or [Ok U64, Err []] which meant the function layout was U64 since there was functionally a single tag in the Result union.

Attempting to enable type unification with early returns now.

[smores56]

This has been fixed and is ready for review, pending tests working in CI as they did on my local machine.