Lambda Return Type Inference

Question

Writing my own JVM compiler, I am facing a giant problem that I am desperately unable to solve:

Lambda Return Type Inference

1. Overview of the compiler lifecycle

More specifically, the order in which things should happen in a compiler to make this work correctly. First of all, as probably every compiler, mine is divided in several phases that run on the AST and do the work. After tokenizing, parsing and building the ast, the following phases happen in that order:

1. RESOLVE_TYPES - Links type tokens to their classes
2. RESOLVE - Links fields and methods
3. CHECK_TYPES - Re-checks types (boxing, ...), Infers generic type arguments
4. CHECK - More uninteresting semantic checks that could be skipped in a perfect world

The phases in question are RESOLVE and CHECK_TYPES. The first resolves method calls and field accesses / assignments as well as their callees and arguments, basically everything important. Between the two, something happens that I call withType. Basically, the context in which an expression appears calls the withType method of the expression with the type it expects the expression to have.

Example

Consider this method and an invocation:

class List[E]
{
    public List[U] mapped(Function[E, U] mapper) = ...
}

List[String] list = [ "a", "b", "c" ]
var list2 = list.mapped { s => s.toUpperCase }

What the compiler does

This is what is going to happen in my compiler now (we ignore what happens with the List class):

1. Read File, Tokenize, Parse - nothing spectacular
2. Resolve Types - nothing special, link all explicit types (List[String])
3. Resolve - link method calls

It is known that list has the type List[String], so we can search that for methods called mapped. Iterate through all methods in the class, and if the name matches the method name, continue with argument types. This works like this: The argument type is Function[E, U], so call isType(Function[E, U]) on the argument { s => s.toUpperCase }. The Statement List { } delegates the call to the lambda expression, which checks if the type has a functional method.

interface Function[T, R]
{
    R apply(T par1) // <-- Here it is
}

Now, link the method and the type Function[E, U] to the lambda expression.

3. Check Types - Infer the type arguments of generic types

We start with Function[E, U]. It is known that E is a type parameter of the generic type List[E], which is also the type of the callee, but more concrete: List[String]. That means that E can be inferred to String. But here comes the problem: Since we need E for the type of the lambda parameter s, it was not possible to resolve the lambda body, i.e. s.toUpperCase. Because of that, it is not possible to infer U at the moment, so the compiler continues without inferring the right type (it simply uses any, which is java.lang.Object). That means that our inferred type is now Function[String, any].

In the next step, the compiler calls the withType method of the lambda expression, which updates the type stored in the lambda AST node. Now we can actually call lambda.checkTypes, which will a) infer the type of s to String and b) call resolve (!), withType and checkTypes on the lambda body, allowing us to compute the return type of the lambda.

2. The problem

As you can see, although we now have the type of the lambda parameter s and the return type of the lambda (String), the type of the lambda is still Function[String, any], which means the result of the call to mapped and therefore the type of the variable list2 is List[any]. However, since the return type of the lambda is actually String, this is not the correct type (since List is invariant on E, the types are not even compatible). The generic type system of the compiler is not advanced enough to make things like

List[String] list2 = list.mapped { s => s.toUpperCase }

impossible, but neither variant is technically correct as long as the lambda type gets inferred to Function[String, any] while it is actually Function[String, String].

Furthermore, consider this expression:

int length = list.mapped(s => s.replace("a", "")).getFirst.length

The call to mapped returns a List[any], therefore getFirst returns an any, which does not have a length member. Normally, one would expect this expression to work, and in fact lambda expressions work like this in Java and Scala (which makes things a little simpler by enforcing their built-in function types, disallowing FunctionInterfaces).

3. The question

Now that you know the details of my compiler lifecycle, type inference and lambda expressions in my programming language, is there any way I could restructure it to make lambda return type inference work correctly? Although they have very different compiler phases and -lifecycles, Java 8 and Scala have to do similar things after all as well.

4. My potential solution

I was able to come up with at least a sort of solution that covers around 10% of the problem, which would be to infer types in RESOLVE so the lambda can actually call resolve of its body in resolve instead of checkTypes. That means that the return type of the lambda is known after resolve and further calls in the chain have the proper callee type (e.g., the getFirst method). However, I am unsure if that would be possible with my current type inference system. Any opinion / advice is greatly appreciated.

Compiler Source Code:

• GitHub Repository
• IType class
• GenericType class
• Method class (inferType around line 880)
• IValue class
• LambdaExpression class

View at own risk, 90k lines of code.

Answer 1

Your problem is that you link method calls before the program is fully type-checked. In a complicated language with subtype polymorphism and parametric polymorphism, you cannot do any kind of linking before you know the type (or in the case of subtyping, the type bounds) of an expression. The solution is to introduce type variables for expressions of (initially) unknown type, and properly unify type variables with concrete types. If you cannot resolve a type variable, that's an error and not an invitation to deduce the most general type.

A proper type inferer would look at the code

var list2 = list.mapped { s => s.toUpperCase }

in the static environment

list : List[String]
List[E]#mapped[U] : (E => U) => List[U]
String#toUpperCase : () => String

like this:

• infer var list2 = list.mapped { s => s.toUpperCase }
  • infer list.mapped { s => s.toUpperCase }
    • infer list
      • known to be List[String]
      • save E = String to the static environment for this expression.
    • resolves to List[E]#mapped[U]
    • introduce variable $U for U.
    • known to be (E => $U) => List[$U]
    • unify arguments: (E => $U) >: typeof({ s => s.toUpperCase })
      • infer { s => s.toUpperCase }
        • introduce variables $a, $b
        • known to be $a => $b
        • unify (E => $U) >: ($a => $b): E <: $a and $b <: $U. Technically, these are just type bounds, but let's consider them equal.
        • save E = $a, $b = $U to the static environment for this expression
        • save s : $a to the static environment for this expression
        • infer s.toUpperCase
          • infer s
            • known to be String
          • infer .toUpperCase.
            • matches String#toUpperCase
            • known to be () => String
          • is String
        • unify String <: $b, but lets consider them equal
        • save $b = String to the static environment for this expression
      • is ($a => $b) = (String => String)
      • already unified (E => $U) >: ($a => $b).
    • is List[$U] = List[String]
  • infer var list2
    • introduce variable $c
    • save list2 : $c into the static environment for this scope
    • is $c
  • unify $c >: List[String], but let's consider them equal
  • save $c = List[String] into the static environment

After this round of type inference, asking the static environment for the type of list2 gives us List[String]. The above steps roughly follow along the type inference steps for a Hindley-Milner type system, but note that you are probably dealing with subtyping, which makes type inference far more complicated: unifying a type only provides us bounds on the type and very rarely a concrete type. In the above example, I ignored these details and always unified the types as equals.

The second major problem is infering the type of the object on which to dispatch for a method call. In this example, I always tried to infer the type of the object before resolving the method call. However, this is not necessary. A method List[E]#mapper[U] : (E => U) => List[U] could also be viewed as a free function .mapper[E, U]: (List[E], E => U) => List[U]. In other words, the implicit this argument is transformed to an explicit argument. This allows us use type variables for the invocant variable, and makes it easier to deal with a set of overloads. Consider the following example:

static env:
  .add[E]    : (List[E], E) => void
  .add[E]    : (List[E], List[E]) => void
  .add[K, V] : (Map[K, V], K, V) => void
  .add       : (Integer, Integer) => Integer
  a          : $a
  b          : $b

expression:
    a.add(b)

when type-infering that expression, we initially only know that this is some kind of .add call, but we do not know which one – we cannot resolve it at this point. We have to unify each of the possible types with the given types. If one unification fails, we remove that possibility from the set of overloads.

• infer a.add(b)
  • candidate (List[E], E) => void
    • variable $c for E
    • unify arguments
      • $a = List[$c]
      • $b = $c
    • is void
  • candidate (List[E], List[E]) => void
    • variable $c for E
    • unify arguments
      • $a = List[$c]
      • $b = List[$c]
    • is void
  • candidate (Map[K, V], K, V) => void
    • variable $c for K
    • variable $d for V
    • unify arguments
      • $a = Map[$c, $d]
      • $b = $c
      • ERROR: no argument for V-parameter
    • discard candidate
  • candidate (Integer, Integer) => Integer
    • unifiy arguments
      • $a = Integer
      • $b = Integer
    • is Integer

So after doing a round of type inference, we are still left with a number of choices. Some languages employ some ranking mechanism, e.g. methods on subtypes are preferred over methods on supertypes. In our example, there is no sensible ranking. If you want your type-checker to be moderately well-performing, you should issue a compiler exception here, e.g.

foo.sourcefile:42:3 error: could not resolve call to method `add`
    a.add(b)
      ^
    candidates are:
    List[E]#add(E elem)
    List[E]#add(List[E] elems)
    Integer#add(Integer rhs)

Likewise, you should display an error when the set of possible methods is empty after attempting unification.

Note that for a different static environment, the the set of possible solutions would only contain one method which you could bind to. E.g.:

a : List[$c]
b : $c

Since the call can be resolved unambiguously in that case, we can continue doing type inference for the rest of the compilation unit in hopes of resolving the $c type variable.

Depending on how you implement generics, you could link the method call directly after successfully infering the type for the call. If you need to know the value of the type variable $c, you need to wait until type inference for the compilation unit has completed, and can do linking in a second pass. This is now easy to do since all expressions are already annotated with their inferred type.