Skip to content

Instantly share code, notes, and snippets.

@jake-87
Last active February 27, 2024 03:09
Show Gist options
  • Save jake-87/1d5a21c2bbe2948a89a6f4d151de74c6 to your computer and use it in GitHub Desktop.
Save jake-87/1d5a21c2bbe2948a89a6f4d151de74c6 to your computer and use it in GitHub Desktop.

A somewhat brief overview of typechecking System F

This document comprises an example of typechecking a type system based on System F, in a small ML-like language.

You can skip the below section if you already understand System F itself.

A brief overview of System F

System F is the name given to an extension of the simply typed lambda calclus.

The expression syntax can be described as the following:

Abstraction: λ(<var> : <type>) . <expr>
Application: <expr> <expr>
Type Abstraction: Λ<var> . <expr>
Type Application: <expr> [<type>]

And of course, you can add your types of choice such as bools or integers.

Here's an example of what the classic id looks like:

id = ΛT. λ(x : T). x

In words, id takes a type as it's first argument, then a variable of that type, and then returns said variable. Here's us calling id with a boolean and an integer:

id [Bool] true
id [Int] 5

Returning exactly what you'd expect.

Now, how do we assign types to these, and what do those types represent?

We can think of Λ, our "type abstraction", or, as it will be referred to in the rest of this document, a "type lambda", as an "introduction rule" for a forall in a type. Likewise, our regular abstraction, or "lambda", can be thought of as an introduction a function arrow.

When we wish to assign the below a type, we follow this process:

id : ???
id = ΛT. λ(x : T). x

First we encounter a type lambda, so we know that the first part of the type will be a forall:

id : forall T, ???
id = ΛT. λ(x : T). x

Then our lambda, which introduces an arrow, and because it's of type T we know the left hand side will be T:

id : forall T, T -> ???
id = ΛT. λ(x : T). x

And then of course we're returning something of type T, so we get:

id : forall T, T -> T
id = ΛT. λ(x : T). x

When we call id like id [Int] 5, the first type application "removes" a forall, and takes forall T, T -> T to Int -> Int, which we can then call with 5.

In this way, like how application is the "eliminator" of lambda abstraction, type application is the "eliminator" of type abstraction.

So far, this isn't anything out of the ordinary for a polymorphic lambda calculus. However, the following demonstrates what a regular polymorphic lambda calculus cannot do:

// we assume some extension of System F with tuples
app2 : forall A B, (forall C, C -> C) -> A -> B -> (A, B)
app2 = ΛΑ. ΛΒ. λ(f : forall C, C -> C). λ(a : A). λ(b : B). (f [A] a, f [B] b)

Let's break this down. First, matching the first two foralls in the type signature, we take two types. Then we take a function of type forall C, C -> C, and call it f.

This is a higher ranked type, and it's what makes System F special. You can use a forall anywhere you want - no restrictions.

We then take our two arguments of type A and B respectively. Then, we call f on both a and b, putting our type applications in - remember, type application serves to "instantiate" our polymorphic type forall C, C -> C and turns it into A -> A and B -> B respectively.

Our example language

All code examples are going to be in psudo-OCaml. It should run with minimal errors if pasted in a file, and rec/and being substituted in some places.

This is the little language we'll be typechecking:

type ty =
  | Name of string
  | Arrow of ty * ty
  | Tup of ty * ty
  | Forall of string * ty

So for example, the following is a valid type: forall T, T -> (T, T).

Here is the expression syntax:

type expr =
  | Bool of bool
  | Int of int
  | Var of string
  | LetIn of string * expr * expr * expr
  | App of expr * expr
  | IfThen of expr * expr * expr
  | Tuple of expr * expr
  | Lambda of string * ty * expr
  | TypeLambda of string * expr

The following is a valid expression:

let x = true in
if x then 5 else 10

Note:

  • We do not have type application in the given syntax. This is very deliberate - given enough information, we can fully infer it.
  • All lambdas must be given a type; inference is out of the scope of this example.

Finally, we have the notion of some "toplevel" expression:

type toplevel = string * ty * expr

We don't include arguments; you can use lambdas.

We'll be using bidirectional typing for the example. This means that we have a check function, that checks a term against a type, and an infer function, that attempts to infer the type of a term. context represents some imaginary (for the purposes of this example) context that stores variable and toplevel types. A trivial one might be:

type context = {
  vars: (string * ty) list; (* mapping from variable name -> type *)
  typs: string list; (* all types that are currently in context (e.g. int) *)
}
let check (ctx: context) (term: expr) (typ: ty): unit = ...
let infer (ctx: context) (term: expr): ty = ...

Let's start by building out the basics of each:

let check (ctx: context) (term: expr) (typ: ty): unit =
  match term, typ with
  | Bool b, Name "bool" -> ()
  | Int i, Name "int" -> ()
  | TypeLambda (name, expr), Forall(tyname, rest) ->
    let ctx_with_ty = add_type_to_context ctx name in
    let rest_with_sub = subst_name rest ~old:tyname ~new:name in
    check ctx_with_ty expr rest_with_sub

In the TypeLambda case, let's say we're checking

check:
ΛT. λ(x : T). x
against
forall a, a -> a

We want to first add T to our context of types, and then subtitute in T everywhere a appears in the rest of the type signature, takng us from

check:
ΛT. λ(x : T). x
against
forall a, a -> a

to

check:
λ(x : T). x
T -> T

subst_name is fairly trivial to implement - just keep in mind alpha equivance, and ensure you don't accidentlly subtitute something you're not meant to.

We continue with the implementation of check:

  | IfThen (scrut, iftrue, iffalse), ty ->
    check ctx scrut (Name "bool");
    check ctx iftrue ty;
    check ctx iffalse ty (* Both branches must have the same type *)
  | Tuple(left, right), Tup(l, r) ->
    check ctx left l;
    check ctx right r (* Check each side respectively *)
  | Lambda (name, ty, expr), Arrow(input, output) ->
    unify ty input;
    let ctx_with = add_var ctx name ty in
    check ctx_with expr output

Here we see the first occurence of unify - this will show up later, but for now, just think of it as checking two types are equal.

Then, application and the default case:

  | App(fn, arg), ty -> failwith "TODO"
  | tm, expected ->
    let inferred = infer ctx tm in
    unify inferred expected

Note that we have left out application for right now - this is because we will come back and do inference for type arguments later.

Now for infer:

let infer (ctx: context) (term: expr): ty =
  match term with
  | Bool b -> Name "bool"
  | Int i -> Name "int"
  | Var v -> lookup_var_type ctx v
  | LetIn(name, this, that) ->
    let thisType = infer ctx type in
    let ctx_with = add_var ctx name thisType in
    infer ctx_with that
  | IfThen(scrut, iftrue, iffalse) ->
    check ctx scrut (Name "bool");
    let iftrueType = infer ctx iftrue in
    check ctx iffalse ifTrueType
  | Tuple(a, b) ->
    let aType = infer ctx a in 
    let bType = infer ctx b in
    Tup(aType, bType)
  | Lambda(name, ty, body) ->
    let ctx_with = add_var ctx name ty in
    let bodyType = infer ctx_with body in
    Arrow(ty, bodyType)
  | TypeLambda(name, body) ->
    let ctx_with = add_type ctx name in
    let bodyType = infer ctx_with body in
    Forall(name, bodyType)

Once again we leave out the application case, and the body of unify. They will be clarified below.

How to infer type application

Let's imagine we're calling id with the argument 5. How do we infer that we want to apply the type Int first? One method is with metavariables. These are best described as "concrete, but unknown", i.e. you don't know what they are yet, but they can't have more then one type.

Walking through how we might use these to infer the result of id 5:

id : forall a, a -> a
id 5 : ???

Step 1: Instiantiate the top level foralls (i.e. any on the most outside part of the type) with metavariables, which we will write ?0, ?1, etc. Also note that while we will write id : ..., this transform is only local in the application.

id : ?0 -> ?0
id 5 : ???

Then when we apply something, we unify these types, solving for metavariables. Unifing the left hand side of id's type with the type of 5:

unify ?0 and (infer ctx 5)
=>
unify ?0 and Int

We utilize shared references for these metavariables, such that when they are solved in one place, they are solved globally, which then gives:

id : Int -> Int
id 5 : Int

For a more complex example, let's walk through how we'd typecheck the first System F example given, rewritten in our little language:

app2 : forall A B, (forall C, C -> C) -> A -> B -> (A , B)
app2 =
  lam f => lam a => lam b =>
  (f a, f b)

We get to (f a, f b) with the following bindings:

types in ctx = ["A", "B"]
bindings in ctx = [
  f : forall C, C -> C
  a : A
  b : B
]

Now we check each side:

check (f a) against A
instantiate `f` with metavariables, given `?1 -> ?1`
unify the left hand side of `f` with the type of `a`, giving `A -> A`
take the right hand side, as that's the return type, giving `A`
unify (f a) and A
==
unify A and A
==
true

Applying this to our model

In order to utilize this inference method, we need to add a case for metavariables to our ty construct, a case for "bound foralls", which will be explained later, a new metavariable type, and a few new functions.

type ty =
  | Name of string
  | Arrow of ty * ty
  | Tup of ty * ty
  | Forall of string * ty
  | Meta of metavar ref
  | BoundForall of metavar ref * ty

and metavar =
  | Unsolved of int
  | Solved of ty

In short, BoundForall is for when instantiation takes place. We want to instantiate something like forall A, A -> A to ?0 -> ?0, but in some cases we have to turn it back again, and therefore we need to know where the forall was. A BoundForall is not used except for in this case.

We add a metavar type that can either be unsolved with a tag, or solved with a type. We then add the following functions:

let fresh_int =
  let i = ref 0 in
  fun () -> 
    incr i;
    !i

let new_metavar () =
  let i = fresh_int () in
  Unsolved i

fresh_int is simply a way to get a new integer each time the function is called.

let subst_name_for_ty (substee: ty) (name: string) (new_ty: ty): ty = ... 

This function functions similiarly to subst_name, but replaces it with a full type instead.

The most important function we add is the below, reify.

let reify (typ: ty) =
  match typ with
  | Name _ -> ty
  | Tup (a, b) -> Tup (reify a, reify b)
  | Arrow(a, b) -> Arrow (reify a, reify b)
  | Forall(s, body) -> Forall(a, reify body)
  | BoundForall (m, body) ->
    begin match m with
      | Solved _ -> reify body (* If it's been solved, we don't need to track it anymore *)
      | Unsolved i -> BoundForall(m, reify body)
    end
  | Meta m ->
    match m with
    | Unsolved i -> Meta (Unsolved i)
    | Solved ty -> reify ty

This takes something with solved metavariables in it, and removes the solved ones, leaving behind only the unsolved one and the default type.

Another function we need is the function that performs instantiation:

let instantiate_top (typ: ty): ty =
  match typ with
  | Forall(s, body) ->
    let m = new_metavar () in
    let body' = instantiate_top body in
    BoundForall(m, subst_name_for_ty body' s m)
  | _ -> typ

Note that we stop instantiating as soon as we hit something that isn't a Forall.

Now let's write (part of) unify, and you can see how we solve metavariables.

let unify (t1: ty) (t2: ty): unit = 
  match reify t1, reify t2 with
  | Name a, Name b -> assert (a == b)
  | Arrow(fn1, arg1), Arrow(fn2, arg2) ->
    unify fn1 fn2;
    unify arg1 arg2
  | Tup(a1, b1), Tup(a2, b2) ->
    unify a1 a2;
    unify b1 b2
  | Forall(name1, body1), Forall(name2, body2) ->
    let body2Inst = subst_name body2 ~old:name2 ~new:name1 in
    (* Here we transform the second body in order to be able to unify easier. Remember to respect alpha! *)
    unify body1 body2Inst
  | BoundForall(_, body), ty
  | ty, BoundForall(_, body) -> unify ty body
    (* We ignore the bound foralls *)
  | Meta m, ty
  | ty, Meta m ->
    begin match m with
      | Solved _ -> failwith "reify should have removed solved metas"
      | Unsolved i -> m := Solved ty
    end
  | _, _ -> failwith "can't unify"

The magic happens in the Meta case. If it's solved, and we have some type, we set the meta (through a reference) to the type. This is what allows us to solve Int applied to ?0 -> ?0 as Int -> Int.

Therefore, our App case in check becomes:

  | App(fn, arg), ty ->
    let fnInst = instantiate_top fn in
    let argType = infer ctx arg in
    begin match fnInst with
      | Arrow(a, b) ->
        unify a argType;
        unify b ty
       | _ -> failwith "Can't apply to a non-arrow type"
    end

And in infer becomes:

  | App(fn, arg) ->
    let fnInst = instantiate_top fn in
    let argType = infer ctx arg in
     begin match fnInst with
      | Arrow(a, b) ->
        unify a argType;
        b
      | _ -> failwith "Can't apply to a non-arrow type"
    end

Now, for the moment, we have it pretty good - we can infer type arguments nicely utilizing our metavars. If you have a good memory, you may remember that we have a BoundForall construct that we're not currently using - here's where it comes in.

For anyone that has implemented Hindley-Milner inference before, let-generalization may be familiar. This is when you're using metavars to infer something, but you forget to remove them when you're binding something in a let block, and the type becomes less general. We don't want our application to affect a variable's usage later.

This means that we need another function to remove these bound foralls, replacing them with normal foralls, and another subst function to replace unsolved metavariables with the correct forall variable. This is what we carry around the integer in Unsolved i for - we can replace each Unsolved i with Name (string_of_int i) and each BoundForall(m, body) where m = ref (Unsolved i) with Forall(string_of_int i, body) (assuming the transform has also been done on body).

We rewrite our LetIn case in infer:

  | LetIn(name, this, that) ->
    let thisType = infer ctx type in
    let generalized = remove_unsolved_metas thisType in
    let ctx_with = add_var ctx name generalized in
    infer ctx_with that

And then you're practically done! Add some toplevel function management, perhaps some lambda inference using similar techniques (challenge: what would you need to do to infer a lambda type if metavars are already in place?), perhaps something fancy like ADTs, and you have a very nice little System F based language.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment