Type inference for Haskell, part 10

This is part 10 of the series on type inference for Haskell.

This post kicks off the middle third of the series. In the first third, we built up a basic type inference system for a subset of Haskell. In this third, each post will add an extension to the type system that has already been created. Today’s extension is something called “let polymorphism.” Once this is added, the type system will be equivalent to Algorithm W (aka Hindley-Milner) defined in the famous Milner-Damas paper, Principal type-schemes for functional programs.

To start off with, let’s look at a motivating example showing why let polymorphism is useful. After that we’ll be ready to define what the language extension actually does. Next will be an example of how a naive implementation might go wrong. To finish up, the simple type inference system will be extended to handle this feature.

A previous post used this example:

let identity = (\x -> x)
in identity identity

As it is right now, our type system cannot handle expressions like these because identity is used as two different types. On the left, it’s (t -> t) -> (t -> t), and on the right it’s just t -> t. As it stands right now, we require a variable to be used as the same type everywhere. You can probably see that this expression would work at runtime, it’s just our type system that cannot handle it.

Let polymorphism solves this issue. It allows a variable declared in a let binding to be used as different types in each usage. The only requirement is that each type it is used as agree with the definition. Exactly what it means to agree with the definition is somewhat subtle.

Before getting in to that, a word on what variables this applies to. Recall how there are two things that add variables to the current scope, let bindings and function arguments. The kind of polymorphism that we are looking at only applies to let bindings (hence the name “let-polymorphism”). Haskell can support polymorphic function arguments with a language extension, but this series isn’t covering language extensions. Google “Rank-N Types” if you want to learn about this extension.

OK, back to how let polymorphism actually works. Each usage of a variable declared in a let-expression can be used as a different type. That type can differ only by certain type variables. In some types, like Int -> Bool, there are no type variables, so there’s nothing to change at each usage. A variable of that type would have to be used as the same type everywhere. In a type like Int -> [a] -> [a], it might be the case that a is a type variable that can be used as different types in each usage.

This process is called instantiation. You might say “the type Int -> [a] -> [a] is instantiated to Int -> [String] -> [String]”. That type could also be instantiated as Int -> [Bool] -> [Bool] at a different usage of the relevant variable. This assumes, of course, that a is a type variable that we are allowed to change at each usage of the variable.

Which Type Variables?

Deciding which type variables are allowed to be changed at each location is the only tricky part of this whole system. Let’s start with an example of where a type variable isn’t allowed to be different at each usage, and then work towards a general rule. Think about the type of f in this expression. From the definition, we don’t know the type of either y or x, so the type will have to be something like ty -> tx.

(\x -> let f = (\y -> x)
       in  [[f 1], f 2])

If you try to find the type of this whole expression in GHCi, you’ll get the error “cannot construct the infinite type.” (This is an error we’ve seen before, but that was for a different reason.) What’s the problem with it? Look at the two calls to f in the list [[f 1], f 2] and note that the first call to f is in a nested list, and the second isn’t. This means that the second call to f needs to return some type that is a list of whatever type the first call to f returned. (Say f 1 returns an Int. That would make [f 1] be an [Int], so f 2 must also return an [Int].) That cannot possibly be the case because both calls to f return exactly the same value (whatever x is bound to).

The problem came from the fact that part of the type of f is determined by something other than the definition of f itself. Though the type of x might still be a type variable, it isn’t one that is actually flexible. (On the other hand, the type of y is one that we can safely change at each usage.) Now we can write the full set of requirements for being able to instantiate part of a type:

• It must be a type variable.
• That type variable must not appear in the type of any of the other bindings in scope.

The reason that the identity example above worked is that the type variable in identity’s type met both of the two requirements.

The code

(THIH Section 8, Type Schemes)

We’ll need some way to represent a type with some type variables that can be replaced. To do this, the Type data structure needs to be extended to handle one more variant, TGen Int.

data Type
  = TVar Tyvar
  | TCon Tycon
  | TAp Type Type
  | TGen Int
  deriving (Eq, Ord, Show)

The TGen variant is used to represent a type variable that can be safely changed at each usage of the variable. Types that must be the same at each usage (such as the type of function arguments) won’t have any of these. The Int field is used to keep track of which of the replaceable type variables this is¹.

We’ll also wrap types in a new structure called a Scheme. A scheme represents a type that is safe to instantiate (though the type might not actually have any variables to replace). The Int field keeps track of how many different replaceable type variables exist in the type. That saves a little work each time the type is instantiated.

data Scheme
  = Forall Int Type
  deriving (Eq, Show)

That structure also gets an instance of the Types typeclass:

instance Types Scheme where
  apply s (Forall n t) = Forall n (apply s t)
  tv (Forall _ t) = tv t

And finally, the Assump (assumption - this stores the variables that are in scope) structure gets updated to use it. We also need to update the functions that use Assump:

data Assump = Id :>: Scheme

instance Types Assump where
  apply s (i :>: sch) = i :>: (apply s sch)
  tv (_ :>: sch)      = tv sch

find :: Monad m => Id -> [Assump] -> m Scheme
find i []              = fail ("unbound identifier: " ++ i)
find i ((i':>:sch):as) = if i==i' then return sch else find i as

We need a way to create type schemes. The toScheme method is used for function arguments. It converts a type to a scheme without allowing any of the type variables to be instantiated. On the other hand, quantify is used for variables introduced in let block. It might, depending on the list of type variables passed in, make some parts of the type into TGen values. The work done by quantify is sometimes called quantification or generalization.

toScheme :: Type -> Scheme
toScheme t = Forall 0 t

quantify :: [Tyvar] -> Type -> Scheme
quantify vs t = Forall (length s) (apply s t)
 where vs' = [ v | v <- tv t, v `elem` vs ]
       s   = zip vs' (map TGen [0..])

In quantify, the variable vs' is a list of the type variables that appear in both the type t and in the list of variables passed in, vs. The variable s is a substitution from the variables in vs' to some TGen values (with numbers starting at 0)². That substitution is applied to the type t to replace the type variables with the TGen values.

If you gave quantify a type like a -> a and the list of type variables that contained a, it would output a scheme like

Forall 1 (TAp (TAp (TCon $ Tycon "->") (TGen 0)) (TGen 0))

We also need a function to instantiate a type scheme. This is very nearly the reverse of quantify - we go from TGen values back to type variables. The only difference is that the type variables we put back in are totally unique. That lets us change those type variable to any type we want.

freshInst :: Scheme -> TI Type
freshInst (Forall n t) = do
  ts <- mapM (\_ -> newTVar) [1..n]
  return (inst ts t)

class Instantiate t where
  inst  :: [Type] -> t -> t

instance Instantiate Type where
  inst ts (TAp l r) = TAp (inst ts l) (inst ts r)
  inst ts (TGen n)  = ts !! n
  inst _  t         = t

The inst function is analogous to apply. Instead of replacing type variables with types, it replaces TGen values with types. In the line inst ts (TGen n) = ts !! n you can see that it is looking up the nth type in the list passed in. (The Instantiate typeclass is going to be useful in later extensions.)

freshInst creates a list of n new type variables (where n is the number of different variables that had originally been converted to TGen - see quantify) and passes them to inst. The resulting type then has a new type variables where it previously had TGen values.

Now it’s finally time to update the type inference function. We’ll have to update both places where a variable is used and places where a variable is introduced. Type inference for literals and function applications doesn’t change at all.

The change for variables is minor. We just need to call freshInst to convert the scheme to a type:

tiExpr as (Var i)            = do
  sc <- find i as
  t <- freshInst sc
  return t

The logic for constructors changes in the same way, adding a call to freshInst:

tiExpr _  (Const (_ :>: sc)) = do
  t <- freshInst sc
  return t

Abstractions now need to convert the argument’s type to a scheme before adding it to the assumptions:

tiExpr as (Abs i e)          = do
  t <- newTVar
  let sc = toScheme t
  te <- tiExpr ((i :>: sc) : as) e
  return $ t `fn` te

And finally, the real heart of the logic runs when dealing with a let binding:

tiBinding :: Infer Binding [Assump]
tiBinding as (Binding name expr) = do
  t <- tiExpr as expr
  sub <- getSubst
  let t' = apply sub t

  let typeVars = tv t' \\ tv (apply sub as)
  let sc = quantify typeVars t'

  return [name :>: sc]

Recall that this function used to be just

tiBinding :: Infer Binding [Assump]
tiBinding as (Binding name expr) = do
  t <- tiExpr as expr
  return [name :>: t]

The tv t' \\ tv (apply sub as) expression is where it finds the list of type variables that are in the variable’s type and not in the type of any other variable in scope. The current substitution has to be applied to both the type and the assumptions to make sure we have the right set of type variables in both.

Note that the mgu function didn’t get updated. When it is called, all TGen values should have already been removed from the type via instantiation.

Behavior

When a variable’s type is generalized into a scheme, some of the type variables might be ones that can be changed at each usage, and some might not. When that variable gets used somewhere, that usage might restrict the type of the variable in some way. If that restriction applies to one of the type variables that can be used as a different type everywhere, then this won’t impact the type of the other usages of the variable. However, for type variables that aren’t allowed to change at each usage, a restriction will still apply to everywhere that variable is used.

As mentioned at the start, this algorithm is now just as powerful as Algorithm-W, though there are some minor changes in how it is implemented.

The next post will look at how to handle self-referential bindings (such as recursive functions).