Enumerating linear inhabitants

Preliminaries

We start with some imports. We’ll be using the State and Logic monads, as well as some Sets to represent contexts.

> {-# LANGUAGE GeneralizedNewtypeDeriving #-}
> 
> import Control.Applicative ((<$>))
> 
> import Control.Monad
> import Control.Monad.State
> import Control.Monad.Logic
> 
> import qualified Data.Set as S
> import qualified Data.Foldable as F
> import Data.Char

To start out, I’m going to simplify things in several ways (more general versions perhaps to follow in some future posts). First, we restrict ourselves to type schemes, that is, polymorphic types where all the quantifiers are at the very beginning. In other words, we rule out higher-rank types. This simplifies things because we won’t have to explicitly represent quantifiers; every type variable will be implicitly quantified. (Also, higher-rank types are scary.) This is not a huge restriction, however, since most of the types we have been considering (with the notable exception of some of gasche’s ideas) have been type schemes.

Second, for now we’ll leave out pairs. This does greatly restrict what we can do; most of the interesting types we’ve explored involved tupling. But we have to start somewhere! I hope to show how to extend this program to pairs in a future post.

Types and Terms

We start out with a representation of types. Due to our simplifications, we have only type variables (represented by Int) and arrow types. We also define a few abbreviations which we can use to help make examples more readable.

> data Ty  = TyVar Int
>          | Arr Ty Ty
>   deriving (Show, Eq, Ord)
> 
> [a, b, c, d] = map TyVar [0 .. 3]
> (-->) = Arr
> infixr 1 -->

We’ll need a few utility functions for manipulating types. result extracts the final result of a type; for example, the result of \((a \to a) \to b \to c\) is \(c\). args extracts a list of the argument types of an arrow type. For example, the arguments of \((a \to a) \to b \to c\) are \([a \to a, b]\).

> result :: Ty -> Ty
> result (Arr _ t2) = result t2
> result t          = t
> 
> args :: Ty -> [Ty]
> args (Arr t1 t2) = t1 : args t2
> args _           = []

And now for terms: we have term variables (again represented by Int), lambda abstraction, and application. (As a side note if you’ve ever played with these sorts of first-order term representations with binding: since we’re not going to do any reduction or substitution, just building terms, we won’t need to be particularly sophisticated about names or alpha-equivalence or anything of that sort.)

> data Tm = Var Int
>         | Lam Int Tm
>         | App Tm Tm
>   deriving (Show, Eq, Ord)

We’ll need one utility function on terms, to build up the nested application resulting from applying a term to a list of arguments. We also implement a quick and dirty pretty-printer for terms.

> apply :: Tm -> [Tm] -> Tm
> apply = foldl App
> 
> ppr :: Tm -> String
> ppr t = ppr' 0 t ""
>   where ppr' _ (Var x)     = pprv x
>         ppr' p (Lam x t)   = showParen (p > 0) $ 
>           showChar '\\' . pprv x . showString ". " . ppr' 1 t
>         ppr' p (App t1 t2) = showParen (p > 1) $
>           ppr' 1 t1 . showChar ' ' . ppr' 2 t2
> 
> pprv :: Int -> ShowS
> pprv = showChar . chr . (+ ord 'a')

Contexts and the Logic monad

As we’re going about our search for linear inhabitants, we’ll need to keep track of a context which contains the variables available for our use, together with their types. We can represent this as a simple set of pairs. We don’t use a Map because we never need to look up variables by name, instead we will be looking at the types in order to pick out an appropriate variable to use.

> type Context = S.Set (Int, Ty)

Now for the declaration of our custom monad, using some GeneralizedNewtypeDeriving goodness:

> newtype M a = M { unM :: StateT Context Logic a }
>   deriving (Functor, Monad, MonadState Context, MonadLogic, MonadPlus)

What’s the Logic monad? Essentially, it lets us write backtracking searches in a very nice way, and even has combinators for doing fair (interleaved) searches among several alternatives, although we won’t need that facility here. We use return to indicate a successful answer and mzero and mplus (from the MonadPlus class) to indicate failure and choice, respectively. Think of it as a very efficient list monad on steroids. (For further reading on MonadPlus, try reading Doug Auclair’s article in Issue 11 of the Monad.Reader; for more on the Logic monad, see Edward Yang’s article in Issue 15.)

What about that StateT Context bit? Well, while we’re searching for terms we need to keep track of the variables which are available to be used (and their types). When we use a variable, we’ll delete it from the context so that it’s no longer available to be used as we continue our search. And at the end of our search we’ll just require the context to be empty, to ensure that no variables go unused.

Why StateT Context Logic as opposed to LogicT (State Context)? Does it matter? Yes, it matters quite a bit! Monad transformers do not, in general, commute. The LogicT (State Context) monad has a single context that gets threaded around the entire backtracking search; changes to the context are permanent. In contrast, with StateT Context Logic, changes to the context get "rolled back" whenever the search backtracks. Put another way, each alternative search path forks off a "local copy" of the state that doesn’t affect other search paths. It’s this second semantics that we want: if we try something that doesn’t end up working out, we want restore the context to the same state it was in before, and try something else. If you’re wondering how I knew which was which, the honest truth is that I just made up a simple example and tried both! But you can also figure it out by staring at types for a bit: "running" a LogicT (State Context) a computation (using runStateT and observeAll) yields something like ([a],s), whereas running a StateT Context Logic a computation yields [(a,s)].

OK, with that out of the way, we have a few utility functions to define. First, a function to run an M computation by starting it in the empty state and observing all possible results:

> runM :: M a -> [a]
> runM = observeAll . flip evalStateT S.empty . unM

Next, some functions for making up new variables that don’t conflict with anything in the context. We don’t need to do anything fancier than this, since the only time we need to make up a new variable is when we are creating a new lambda abstraction, and the only things that could go inside it are the things currently in the context, or more lambda abstractions.

> freshFor :: Context -> Int
> freshFor ctx | S.null ctx  = 0
>              | otherwise   = 1 + F.maximum (S.map fst ctx)
> 
> fresh :: M Int
> fresh = gets freshFor

Here are functions for extending the context with a new variable and its type, and ensuring that a context is empty. Note how we use guard to ensure the context is empty, so that the whole computation will fail if it is not. The effect is that we can simply state "ctxEmpty" at a certain point in order to require the context to be empty. If it is, the search proceeds; if not, the search backtracks to the last choice point and tries something else.

> extendCtx :: Int -> Ty -> M ()
> extendCtx x ty = modify (S.insert (x,ty))
> 
> ctxEmpty :: M ()
> ctxEmpty = get >>= guard . S.null

The last utility function we will need is chooseBinding, which really forms the heart of the search procedure. It has the effect of nondeterministically choosing and deleting one of the bindings from the current context. To see how it works, first note what the locally defined choose function does: given a binding (that is, a pair of a variable and a type), it deletes that binding from the context and returns the binding. So, chooseBinding gets the context, converts it to a list, and maps choose over it, resulting in a list of actions, each of which will "choose" one of the bindings from the context by deleting and returning it. We then use msum (which simply combines all the elements of a list with mplus) to offer all these actions as independent choices.

> chooseBinding :: M (Int, Ty)
> chooseBinding = msum . map choose =<< gets S.toList
>   where choose bind = modify (S.delete bind) >> return bind

Finding inhabitants

And now for the main enumeration procedure. We run the linearInhabitant computation, which gives us all the solutions. linearInhabitant, in turn, calls inhabitant to generate some inhabitant, and then requires the context to be empty.

> linearInhabitants :: Ty -> [Tm]
> linearInhabitants ty = runM linearInhabitant
>   where linearInhabitant = do
>           inh <- inhabitant ty
>           ctxEmpty
>           return inh

And now for inhabitant. Given a type, how do we generate an inhabitant of that type? Doing this turns out to be surprisingly easy, given our custom monad and the simple tools we have built up.

First, every inhabitant of an arrow type must be a lambda abstraction. Actually, this isn’t precisely true – for example, if there is a variable of type \(a \to b\) in the context and we want an inhabitant of \(a \to b\), we could just use that variable. What I really mean is that it doesn’t hurt to always generate a lambda abstraction for arrow types. Essentially it means we will always generate fully eta-expanded terms, e.g. \(\lambda a. f a\) instead of just \(f\). We could "optimize" things so we generate nicely eta-reduced terms when possible, but I leave this extension as an exercise. For now, if we see an arrow type \(\tau_1 \to \tau_2\), we generate a fresh variable \(x\), put it in the context with type \(\tau_1\), and generate the term \(\lambda x. t\) where \(t\) is an inhabitant of \(\tau_2\).

> inhabitant :: Ty -> M Tm
> inhabitant (Arr ty1 ty2) = do
>   x <- fresh
>   extendCtx x ty1
>   Lam x <$> inhabitant ty2

The other case is when we want an inhabitant of a type variable \(a\). At this point you might (as I did at first) start thinking about two cases, one where there is a variable of type \(a\) in the context and one where there isn’t. But it turns out we only need one general rule. The only way to get something of type \(a\) is to find something in the context whose result type is \(a\), and apply it to some arguments of the right types. This neatly encompasses the degenerate case where something of type \(a\) is already in the context: we just apply it to zero arguments! (If you are at all familiar with the Coq proof assistant, this is essentially the apply tactic.)

So, we nondeterministically choose something from the context and require its result to be the type we are looking for. We then try to generate inhabitants for each of its argument types, and apply it to the generated terms. Notice how we use mapM (as opposed to map and msum) to generate inhabitants for the argument types, which sequences all the computations together. This ensures that the changes each computation makes to the context will be seen by the rest, since we don’t want any of the arguments re-using variables that were already used by previous ones.

> inhabitant a = do
>   (f,ty) <- chooseBinding
>   guard $ result ty == a
>   as <- mapM inhabitant (args ty)
>   return $ apply (Var f) as

Examples

That’s all! Let’s try it out.

> printInhabitants :: Ty -> IO ()
> printInhabitants = mapM_ (putStrLn . ppr) . linearInhabitants

First we’ll try generating some identity functions:

*Main> printInhabitants $ a --> a
\a. a
*Main> printInhabitants $ (a --> b) --> (a --> b)
\a. (\b. a b)

What about something with more than one inhabitant?

*Main> let aa = a --> a in printInhabitants $ aa --> aa --> aa --> aa
\a. (\b. (\c. (\d. a (b (c d)))))
\a. (\b. (\c. (\d. a (c (b d)))))
\a. (\b. (\c. (\d. b (a (c d)))))
\a. (\b. (\c. (\d. b (c (a d)))))
\a. (\b. (\c. (\d. c (a (b d)))))
\a. (\b. (\c. (\d. c (b (a d)))))

And here’s something a bit fancier, with some sort of continuation mumbo-jumbo:

*Main> printInhabitants $ (((a --> b) --> b) --> c) --> a --> c
\a. (\b. a (\c. c b))

Cool, it seems to work! Next time I’ll show how to extend this to handle pairs as well (but of course, feel free to save this post into an .lhs file and try extending it to handle pairs yourself).