Disciple-style Regions in Haskell, Part 1

I've admired ddc for quite a while, in particular its region system. Unfortunately, it also has many annoyances, and one big one is complexity. To rectify this, I've created a system implementing regions in Haskell.

The Essence of Regions

At first we might adopt a definition of regions based on their use in memory management: they are areas where you can allocate memory and work with it, as show by the class below.

-- WRONG!
class (Monad (Environment r)) => Region r where
    data Ref r :: * -> *
    type Environment r :: * -> *

    newRef :: a -> Environment r (Ref r a)
    readRef :: Ref r a -> Environment r a
    writeRef :: a -> Ref r a -> Environment r ()

There are a couple of instances of this - consider ST, or IO, or STM. Unfortunately, this is not what we want. We need to support, for example, immutable regions. Not only that, but in rare cases we might want a write only region (as a random example, a password store). You might have regions that you can't create data in. In the end, we have reduced the idea of a region to something very small: a place with data.

data family Ref r :: * -> *

Note that this also allows us to do away with the monad. Of course, mutation and reading and creation are common, so we should have some classes:

class (Monad m) => Writable r m where
    writeRef :: a -> Ref r a -> m ()

class (Monad m) => Readable r m where
    readRef :: Ref r a -> m a

class (Monad m) => Creatable r m where
    newRef :: a -> m (Ref r a)

Following the previous idea of not needing a fixed monad, we do not restrict a reference to have a single monad assosciated with it.

Some Examples

data Mut s
newtype instance Ref (Mut s) a = MutRef (STRef s a)

instance Writable (Mut s) (ST s) where
    writeRef val (MutRef ref) = writeSTRef val ref

instance Readable (Mut s) (ST s) where
    readRef (MutRef ref) = readSTRef ref

instance Creatable (Mut s) (ST s) where
    createRef val = fmap MutRef $ newSTRef val

data Immut
newtype instance Ref Immut a = ImmutRef a

instance (Monad m) => Readable Immut m where
    readRef (ImmutRef val) = return val

instance (Monad m) => Creatable Immut m where
    createRef val = return (ImmutRef val)

data Atomic
newtype instance Ref Atomic a = AtomicRef (TVar a)

instance Writable Atomic STM where
    writeRef val (AtomicRef ref) = writeTVar val ref

instance Readable Atomic STM where
    readRef (AtomicRef ref) = readTVar ref

instance Creatable Atomic STM where
    createRef val = fmap AtomicRef $ newTVar val

instance Writable Atomic IO where
    writeRef val (AtomicRef ref) = writeTVarIO val ref

instance Readable Atomic IO where
    readRef (AtomicRef ref) = readTVarIO ref

instance Creatable Atomic IO where
    createRef val = fmap AtomicRef $ newTVarIO val

To support things like pointers or specialized references, we probably want to add the ability to restrict the types that can be put in a Ref, but this is beyond the scope of this post.

Building Complex Data Structures

As an example, I will build a cyclic doubly linked list with this framework. Actually constructing the representation and basic operations of a complex data structure is pretty easy:

data DList r a = DList (Ref r (DList r a)) a (Ref r (DList r a))

head :: (Monad m) => DList r a -> m a
head (DList _ x _) = return x

tail :: (Readable r m, Writable r m) => DList r a -> m (DList r a)
tail (DList lRef _ rRef) = do
    l@(DList ll xl rl) <- readRef lRef
    r@(DList lr xr rr) <- readRef rRef
    writeRef lr l
    writeRef rl r
    return r

singleton :: (Creatable r m, MonadFix m) => a -> m (DList r a)
singleton x = do rec
    l <- createRef result
    r <- createRef result
    let result = DList l x r
    return result

-- And so on...

However, how do we switch between representations? With arrays, we usually have the freeze and thaw methods. Here, we have a possibly infinite collection of regions to switch between - how can we do it? I'll post my answer in part 2, since this post is getting a bit long.

Stream Fusion for Conduits

Recently I've been looking at Michael Snoyman's conduit package. Now, conduit is a pretty impressive library, but I was skeptical about its performance. If you look at its source, you see almost every function recursively going down the stack of conduit operations and creating a new one. Now only that, but deforestation is scarce, so it is hard to justify constantly rebuilding these trees.

To fix this, the first thing I tried was using the free monad implementation described here. Unfortunately, this has a large problem: anything besides monad operations is not easy. Functions like pipe, which composes pipes, require building a structure and then traversing it.

After trying a number of similar methods (combining church encoding with Scott encoding, etc.), all of which were equally unsuccessful, I came upon the solution: Stream Fusion.

Stream Fusion

Stream Fusion (this paper is very readable, by the way) is a techneque for removing unnecessary list or array traversals by using an inversion of control. The key definitions are as follows:

data Step s a = Skip s | Yield s a | Done
data Stream a = forall s . Stream (s -> Step s a) s

stream :: [a] -> Stream a
stream {-list-} = Stream next {-list-} where
    next [] = Done
    next (a : as) = Yield as a

unstream :: Stream a -> [a]
unstream (Stream next s0) = go s0 where
    go s = case next s of
        Done -> []
        Skip s -> go s
        Yield s a -> a : go s

The library on hackage modifies these functions lightly to force constructor specialization to happen, but this is irrelevant.

Next we have a library full of functions on stream, like so:

mapStream :: (a -> b) -> Stream a -> Stream b
mapStream f (Stream next s) = Stream next' s where
    next' s = case next s of
        Done -> Done
        Yield s' a -> Yield s' (f a)
        Skip s' -> Skip s'

filterStream :: (a -> Bool) -> Stream a -> Stream a
filterStream p (Stream next s) = Stream next' s where
    next' s = case next s of
        Done -> Done
        Yield s' a
            | p a       -> Yield s' a
            | otherwise -> Skip s'
        Skip s' -> Skip s'

The filterStream function here is especially interesting, as it shows how the Skip constructor can be useful - it allows filterStream to be non-recursive. This helps GHC's optimizer immensely.

Next, we implement a list library on top of this stream library:

map :: (a -> b) -> [a] -> [b]
map f = unstream . mapStream f . stream

filter :: (a -> Bool) -> [a] -> [a]
filter p = unstream . filterStream p . stream

Note that if we inline these definitions, we get things like this:

map f . filter g . map h =
    unstream . mapStream f .
    stream . unstream .
    filterStream g .
    stream . unstream .
    mapStream h . stream

Notice how stream . unstream shows up a lot? Well, at least approximately, streams and lists are isomorphic, so we can add this rule to our library:

{-# RULES
"stream/unstream" forall s . stream (unstream s) = s
  #-}

With this simplification, we now have only one stream and one unstream in out pipeline. This means that we now have one tight loop, instead of three, increasing performance.

Applying This to Conduits

So far, our discussion has only applied to lists. How do we do something like this for, say, a conduit? The answer lies near the end of the paper, in section 9.3, wchich brings up the example of a binary tree:

data Tree a b = Leaf a | Branch b (Tree a) (Tree a)

To apply stream fusion to this type, we modify the Step type:

data Step a b s = Leaf_s a | Branch_s b s s | Skip s
data Stream a b = forall s . Stream (s -> Step a b s) s

Moving on to conduits, we have what follows.

data Step l i o u m r s =
    HaveOutput s (m ()) o
  | NeedInput (i -> s) (u -> s)
  | Done r
  | PipeM (m s)
  | Leftover s l
  | Skip s
data Stream l i o u m r = forall s . Stream (s -> Step l i o u m r s) s

The full code is available here.

Performance

As an optimization library, benchmarks are pretty improtant. I've only done preliminary testing, but here are the results. Especially for something so quickly thrown together by someone who doesn't understand the nuances of SpecConstr, I'd say the results are very promising.