Addressing Pieces of State with Profunctors

Attempted segue

Since I first wrote about profunctors there has been quite a bit of activity in the area so I think it's about time I revisited them. I could just carry on from where I left off 5 years ago but there have been so many tutorials on the subject that I think I'll have to assume you've looked at them. My favourite is probably Phil Freeman's Fun with Profunctors. What I intend to do here is solve a practical problem with profunctors.

The problem

Arrows are a nice mechanism for building circuit-like entities in code. In fact, they're quite good for simulating electronic circuits. Many circuits are very much like pieces of functional code. For example an AND gate like this

can be nicely modelled using a pure function: c = a && b. But some components, like flip-flops, have internal state. What comes out of the outputs isn't a simple function of the inputs right now, but depends on what has happened in the past. (Alternatively you can take the view that the inputs and outputs aren't the current values but the complete history of the values.)

We'll use (Hughes) arrows rather than simple functions. For example, one kind of arrow is the Kleisli arrow. For the case of Kleisli arrows built from the state monad, these are essentially functions of type a -> s -> (b, s) where s is our state. We can write these more symmetrically as functions of type (a, s) -> (b, s). We can think of these as "functions" from a to b where the output is allowed to depend on some internal state s. I'll just go ahead and define arrows like this right now.

First the extensions and imports:

> {-# OPTIONS -W #-}
> {-# LANGUAGE Arrows #-}
> {-# LANGUAGE RankNTypes #-}
> {-# LANGUAGE FlexibleInstances #-}


> import Prelude hiding ((.), id)
> import Control.Arrow
> import Control.Category
> import Data.Profunctor
> import Data.Tuple

And now I'll define our stateful circuits. I'm going to make these slightly more general than I described allowing circuits to change the type of their state:

> newtype Circuit s t a b = C { runC :: (a, s) -> (b, t) }


> instance Category (Circuit s s) where
>     id = C id
>     C f . C g = C (f . g)


> instance Arrow (Circuit s s) where
>     arr f = C $ \(a, s) -> (f a, s)
>     first (C g) = C $ \((a, x), s) -> let (b, t) = g (a, s)
>                                       in  ((b, x), t)

This is just a more symmetrical rewrite of the state monad as an arrow. The first method allows us to pass through some extra state, x, untouched.

Now for some circuit components. First the "pure" operations, a multiplier and a negater:

> mul :: Circuit s s (Int, Int) Int
> mul = C $ \((x, y), s) -> (x*y, s)


> neg :: Circuit s s Int Int
> neg = C $ \(x, s) -> (-x, s)

And now some "impure" ones that read and write some registers as well as an accumulator:

> store :: Circuit Int Int Int ()
> store = C $ \(x, _) -> ((), x)


> load :: Circuit Int Int () Int
> load = C $ \((), s) -> (s, s)


> accumulate :: Circuit Int Int Int Int
> accumulate = C $ \(a, s) -> (a, s+a)

I'd like to make a circuit that has lots of these components, each with its own state. I'd like to store all of these bits of state in a larger container. But that means that each of these components needs to have a way to address its own particular substate. That's the problem I'd like to solve.

Practical profunctor optics

In an alternative universe lenses were defined using profunctors. To find out more I recommend Phil Freeman's talk that I linked to above. Most of the next paragraph is just a reminder of what he says in that talk and I'm going to use the bare minimum to do the job I want.

Remember that one of the things lenses allow you to do is this: suppose we have a record s containing a field of type a and another similar enough kind of record t with a field of type b. Among other things, a lens gives a way to take a rule for modifying the a field to a b field and extend it to a way to modify the s record into a t record. So we can think of lenses as giving us functions of type (a -> b) -> (s -> t). Now if p is a profunctor then you can think of p a b as being a bit function-like. Like functions, profunctors typically (kinda, sorta) get used to consume (zero or more) objects of type a and output (zero or more) objects of type b. So it makes sense to ask our lenses to work with these more general objects too, i.e. we'd like to be able to get something of type p a b -> p s t out of a lens. A strong profunctor is one that comes pre-packed with a lens that can do this for the special case where the types s and t are 2-tuples. But you can think of simple records as being syntactic sugar for tuples of fields, so strong profunctors also automatically give us lenses for records. Again, watch Phil's talk for details.

So here is our lens type:

> type Lens s t a b = forall p. Strong p => p a b -> p s t

Here are lenses that mimic the well known ones from Control.Lens:

> _1 :: Lens (a, x) (b, x) a b
> _1 = first'


> _2 :: Lens (x, a) (x, b) a b
> _2 = dimap swap swap . first'

(Remember that dimap is a function to pre- and post- compose a function with two others.)

Arrows are profunctors. So Circuit s s, when wrapped in WrappedArrow, is a profunctor. So now we can directly use the Circuit type with profunctor lenses. This is cool, but it doesn't directly solve our problem. So we're not going to use this fact. We're interested in addressing the state of type s, not the values of type a and b passed through our circuits. In other words, we're interested in the fact that Circuit s t a b is a profunctor in s and t, not a and b. To make this explicit we need a suitable way to permute the arguments to Circuit:

> newtype Flipped p s t a b = F { unF :: p a b s t }

(It was tempting to call that ComedyDoubleAct.)

And now we can define:

> instance Profunctor (Flipped Circuit a b) where
>     lmap f (F (C g)) = F $ C $ \(a, s) -> g (a, f s)
>     rmap f (F (C g)) = F $ C $ \(a, s) -> let (b, t) = g (a, s)
>                                           in  (b, f t)


> instance Strong (Flipped Circuit a b) where
>     first' (F (C g)) = F $ C $ \(a, (s, x)) -> let (b, t) = g (a, s)
>                                                in  (b, (t, x))

Any time we want to use this instance of Profunctor with a Circuit we have to wrap everything with F and unF. The function dimap gives us a convenient way to implement such wrappings.

Let's implement an imaginary circuit with four bits of state in it.

Here is the state:

> data CPU = CPU { _x :: Int, _y :: Int, _z :: Int, _t :: Int } deriving Show

As I don't have a complete profunctor version of a library like Control.Lens with its template Haskell magic I'll set things up by hand. Here's a strong-profunctor-friendly version of the CPU and a useful isomorphism to go with it:

> type ExplodedCPU = (Int, (Int, (Int, Int)))


> explode :: CPU -> ExplodedCPU
> explode (CPU u v w t) = (u, (v, (w, t)))


> implode :: ExplodedCPU -> CPU
> implode (u, (v, (w, t))) = CPU u v w t

And now we need adapters that take lenses for an ExplodedCPU and (1) apply them to a CPU the way Control.Lens would...

> upgrade :: Profunctor p =>
>            (p a a -> p ExplodedCPU ExplodedCPU) ->
>            (p a a -> p CPU CPU)
> upgrade f = dimap explode implode . f


> x, y, z, t :: Flipped Circuit a b Int Int -> Flipped Circuit a b CPU CPU
> x = upgrade _1
> y = upgrade $ _2 . _1
> z = upgrade $ _2 . _2 . _1
> t = upgrade $ _2 . _2 . _2

...and (2) wrap them so they can be used on the flipped profunctor instance of Circuit:

> (!) :: p s t a b -> (Flipped p a b s t -> Flipped p a b s' t') ->
>        p s' t' a b
> x ! f = dimap F unF f x

After all that we can now write a short piece of code that represents our circuit. Notice how we can apply the lenses x, ..., t directly to our components to get them to use the right pieces of state:

> test :: Circuit CPU CPU () ()
> test = proc () -> do
>     a  <- load ! x       -< ()
>     b  <- load ! y       -< ()
>     c  <- mul            -< (a, b)
>     d  <- neg            -< c
>     e  <- accumulate ! t -< d
>     () <- store ! z      -< e


>     returnA              -< ()


> main :: IO ()
> main = do
>     print $ runC test ((), CPU 2 30 400 5000)

Of course with a suitable profunctor lens library you can do a lot more, like work with traversable containers of components.

Note that we could also write a version of all this code using monads instead of arrows. But it's easier to see the symmetry in Flipped Circuit when using arrows, and it also sets the scene for the next thing I want to write about...