Higher-rank type constraints

I recently encountered a Haskell function which apparently requires "higher-rank" type constraints. I worked around the issue with some type trickery. I'd like to share the problem and the workaround, and especially get feedback on whether there's a better solution.

I did say type trickery:

{-# LANGUAGE
    GADTs
  , FlexibleContexts
  , RankNTypes
  , ScopedTypeVariables #-}

Suppose we're implementing an in-place sorting algorithm. We'll use the ST monad, so that we can expose a pure interface. Assuming that we only want to sort primitive types, we will use the unboxed STUArray type, which should improve space and time usage. Actually, it's totally unwarranted premature optimization for the sake of the problem. :)

import Control.Monad.ST
import Data.Array.ST

The monadic part of this code is simple enough. I've left out the part where we actually, you know, sort the list. Consult your bedside copy of CLRS.

sortM
  :: forall a s.
     (Ord a, MArray (STUArray s) a (ST s))
  => [a]
  -> ST s [a]
sortM xs = do
  arr <- newListArray (1, length xs) xs
           :: ST s (STUArray s Int a)
  -- do some in-place sorting here
  getElems arr

Note that sortM has a constraint MArray (STUArray s) a (ST s). All we're really saying is that a is a type primitive enough that STUArray can unbox it. The rest of the constraint just connects the concrete world of STUArray with the generic MArray interface.

Note also that we give an explicit type on our call to newListArray. Otherwise, there's nothing to constrain what sort of array we get, and GHC will complain of ambiguity. Furthermore, we need to unify this a and s with the a and s from the signature on sortM. We do that using the ScopedTypeVariables extension and the explicit forall. (If you turn on UnicodeSyntax you can write it as ∀; how cool is that?)

Good enough so far. Now, the point of using ST was to expose a pure interface, right?

sortP_1 :: (Ord a) => [a] -> [a]
sortP_1 xs = runST (sortM xs)

Suddenly GHC is very unhappy:

    Could not deduce (MArray (STUArray s) a (ST s)) from the context ()
      arising from a use of `sortM' at code.hs:21:20-27
    Possible fix:
      add (MArray (STUArray s) a (ST s)) to the context of
        the polymorphic type `forall s. ST s a'
      or add an instance declaration for (MArray (STUArray s) a (ST s))
    In the first argument of `runST', namely `(sortM xs)'
    In the expression: runST (sortM xs)
    In the definition of `sortP_1': sortP_1 xs = runST (sortM xs)

Okay, fair enough. We've given no reason for GHC to believe that a is an STUArray-compatible type. Let's copy over the constraint from sortM:

sortP_2
  :: (Ord a, MArray (STUArray s) a (ST s))
  => [a] -> [a]
sortP_2 xs = runST (sortM xs)

Still no good:

    Could not deduce (MArray (STUArray s1) a (ST s1))
      from the context ()
      arising from a use of `sortM' at code.hs:27:20-27
    Possible fix:
      add (MArray (STUArray s1) a (ST s1)) to the context of
        the polymorphic type `forall s. ST s a'
      or add an instance declaration for (MArray (STUArray s1) a (ST s1))
    In the first argument of `runST', namely `(sortM xs)'
    In the expression: runST (sortM xs)
    In the definition of `sortP_2': sortP_2 xs = runST (sortM xs)

Note that GHC gives an error about the variable s1, not s. This is a valuable hint. Our caller (notionally) chooses one s when she provides evidence to satisfy the constraint MArray (STUArray s) a (ST s). Then runST (notionally) gets to choose a different s because of its higher-rank type.

What to do? It's useless for our caller to pick some particular s and provide evidence. We need that evidence for all s:

sortP_3
  :: (Ord a, forall s. MArray (STUArray s) a (ST s))
  => [a] -> [a]
sortP_3 xs = runST (sortM xs)

Unfortunately, GHC does not support this "higher-rank constraint":

code.hs:32:13: malformed class assertion

Well... the -> of functions and the => of constraints are not as different as they may appear. A Haskell compiler can implement type classes by translating constraints to explicit "dictionary" arguments. Let's do the same thing ourselves:

data Evidence s e where
  Evidence
    :: (MArray (STUArray s) e (ST s))
    => Evidence s e

We don't quite need to reproduce the MArray methods in our "dictionary". (That's good, because some of them are hidden. A strange class indeed.) It's enough to write a GADT constructor with this constrained type. The constructor will implicitly capture the constraint evidence (such as a dictionary) from its call-site.

Now we need to express the universal quantification on s:

data ElemType e = ElemType (forall s. Evidence s e)

A (non-⊥) value of type ElemType e is evidence that e is a suitable element type for STUArray. That's all we wanted the caller to say in the first place!

Now we can pass the appropriate evidence:

sortP_4
  :: forall a.
     (Ord a)
  => ElemType a
  -> [a] -> [a]
sortP_4 (ElemType Evidence) xs = runST (sortM xs)

Not quite:

    Couldn't match expected type `forall s. Evidence s a'
           against inferred type `Evidence s e'
    In the pattern: Evidence
    In the pattern: ElemType Evidence
    In the definition of `sortP_4':
        sortP_4 (ElemType Evidence) xs = runST (sortM xs)

GHC still wants some coaxing to match up this s with the other s. We push the pattern-matching of Evidence inside runST:

sortP_5
  :: forall a.
     (Ord a)
  => ElemType a
  -> [a] -> [a]
sortP_5 (ElemType e) xs = runST (f e) where
  f :: Evidence s a -> ST s [a]
  f Evidence = sortM xs

And...

[1 of 1] Compiling Main             ( code.hs, interpreted )
Ok, modules loaded: Main.
*Main>

A glorious sight! Then, a moment of concern: is sortP_5 actually usable? Does ElemType have non-⊥ values? Let's check:

*Main> sortP_5 (ElemType Evidence) ([3,4] :: [Int])
[3,4]

Nice. It's sorted and everything. :) As expected, this works only for unboxable types:

*Main> sortP_5 (ElemType Evidence) ["foo","bar"]

<interactive>:1:18:
    Could not deduce (MArray (STUArray s) [Char] (ST s))
      from the context ()
      arising from a use of `Evidence' at <interactive>:1:18-25

As usual, this interesting exercise came from a question on the #haskell IRC channel. That's why it's one of my favorite places, even if they do specialize in making GHC cry.