Type classes in the time of recession

This post is about using GHC’s rewrite rules engine to implement really really cheap type class instance resolution at compile time. Have fun, and remember: Those who would give up essential type safety, to purchase a little temporary liberty, deserve neither liberty nor type safety.

Haskell type classes are a powerful technique for bounded parametric polymorphism. Consider the type class Eq:

    class  Eq a  where
        (==), (/=)           :: a -> a -> Bool
        x /= y               = not (x == y)
        x == y               = not (x /= y)

This defines the (==) and (/=) equality operators, and gives some default implementations (that mean we need only define one of the two to get both functions).

Now, equality only makes sense on a bounded set of types — not all types support it. Now, we can add a type to the set of types that support equality by writing an instance of Eq for it. For example, with Floats, we can define equality in terms of primitive machine equality on floats:

    instance Eq Float where
        (F# x) == (F# y) = x `eqFloat#` y

Or for unbounded integers, in terms of the GMP integer routines:

    instance  Eq Integer  where
        (==) = gmp_eqInteger
        (/=) = gmp_neqInteger

So now we can write:

    main = do
        print (pi == exp 1)
        print (42 == 43-1)

And everything works nicely.

We get to reuse the same symbol, (==), for equality on any types that support it.

Now, type classes desugar to dictionaries, in GHC, passing the instance methods around. That is,

        (==) :: a -> a -> Bool

Becomes:

        (==) :: EqDict a -> a -> a -> Bool

where the dictionary stores the equality function to use when testing ‘a’s for equality. And when the type for ‘a’ is known statically, GHC usually does the right thing, and resolves it statically.

What isn’t widely known is that we can hack our own static dispatch system together in GHC, using the compiler’s term rewriting capabilities — known as rewrite rules. These are described in the paper “Playing by the rules: Rewriting as a practical optimization technique in GHC”.

First, to turn on rewrite rule syntax, enable -fglasgow-exts:

    {-# OPTIONS -fglasgow-exts #-}

Now, we need to define the type class operation we want to overload:

    class  MEEq a  where
        (=-=)   :: a -> a -> Bool

We’ll just define the type here, no concrete default implementation.

The next thing to do is to tell the compiler which types we are going to have support (=-=):

    instance MEEq ()
    instance MEEq Bool
    instance MEEq Int

Now this is a bit evil, since there’s still no concrete implementation of equality on any of these types. So if you tried to use the code at this point it would fail.

Now, we’ll write some standalone implementations of equality on these types. This code would normally go in a instance body, but we’ll just have them float free:

    eq_Bool :: Bool -> Bool -> Bool
    eq_Bool True  True  = True
    eq_Bool False False = True
    eq_Bool _     _     = False

    eq_Unit :: () -> () -> Bool
    eq_Unit () () = True

Ok. So we could write:

    main = print (True `eq_Bool` False)

But that doesn’t let us do any overloading. What we need is a way to spot (=-=) and have the compiler replace it with the appropriate function, based on the method type.

And we can do this. GHC supports rewrite rules, which are used for a wide range of compile time optimizations, including sophisticated data structure transformations.

Normally, rewrite rules are used to match on a name, and replace the left hand side with some new right hand side.

For example,

    {-# RULES
        "map fusion" forall f g.
            map f . map g == map (f . g)
      #-}

Is the rule for map fusion, replacing two traversals of a list with one. A less used feature is the ability to match on an expression’s type on either the left or right hand side of a rewrite rule (the left and right hand sides must have the same type).

This let’s us trivially encode compile-time instance resolution:

    {-# RULES

    "eq/Bool"   (=-=) = eq_Bool :: Bool -> Bool -> Bool

    "eq/Unit"   (=-=) = eq_Unit :: ()   -> ()   -> Bool

      #-}

So, whenever you see (=-=) at Bool or () type, replace it with the concrete implementation. Overloading resolution, statically, for free.

The slight downside is that if the rule fails to fire, there’s no concrete implementation of (=-=) to fall back on at runtime. So you’ll get a runtime failure :-)

We can see this in action with:

    main = do
        print $ True =-= False
        print $ ()   =-= ()
        print $ 7    =-= (8 :: Int)

We can compile this code, and check the optimizer did the right thing with ghc-core:

    $ ghc-core A.hs
    ...
    2 RuleFired
        1 eq/Bool
        1 eq/Unit
    ...

Excellent, the rules fired for Bool and Unit, replacing the abstract (=-=) with its concrete implementation.

Running the resulting code:

    $ ./a.out
    False
    True
    Class.hs:10:0: No instance nor default method for class operation Main.=-=

And all is good, well, except for the Int, which wasn’t resolved, since we didn’t write a rule for it.

We’d like to recover the property that real type classes have — that it’s a type error to use a type class method at a type other than those that support it. Thanks to Gwern Branwen for suggesting we try this.

The problem is: we want to replace any remnant calls to (=-=) with an expression that will fail to compile. However, we’re strongly constrained by the rewrite rules system — rewrite rules must always be syntactically and type correct. However, and this is the deliciously evil part, they don’t need to be confluent or terminating.

So we can encode the missing instance type error as a non-terminating rewrite rule! And then there’s no risk of runtime failure — as the compiler won’t terminate :-)

We do this very simply, with the rule:

    {-# RULES
    "Unable to resolve instance resolution" (=-=) = (=-=)
      #-}

This simply replaces itself with itself forever. Now, our type correct programs work as expected:

    main = do
        print $ True =-= False
        print $ ()   =-= ()

Results in:

    $ ./A 
    False
    True

However, if we introduce an unresolved overloading:

    main = do
        print $ True =-= False
        print $ ()   =-= ()
        print $ 7    =-= (8 :: Int)

We get a “type error”:

    $ ghc -O2 A.hs
    ^C

in the form of the compiler failed to terminate. As all good Haskellers know, this is morally sound: one bottom is as good as another, so a type error is as good as a divergent compiler! And well-typed programs are still not going to go wrong.

We can also encode the type error as a (compile time) linker failure, by replacing (=-=) with an unknown foreign symbol that we import irregardless. Of course, this is very evil.

Rewrite rules are a powerful way for users to extend the optimization suite of the compiler, and they’re made so by purity and static type information, enabling a richer set of possible transformations than would be possible otherwise.