Talking with a friend earlier today we decided to make an experiment in declaring generic mutable variables that can be used in IO, ST or State monads. I don’t think this is useful, but it was fun to write. Here is result.

{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE FunctionalDependencies  #-}

import Data.Array.MArray
import Data.Array.IO
import Control.Monad.ST
import Data.STRef
import qualified Data.List.NonEmpty as NE
import Data.List.NonEmpty (NonEmpty)
import qualified Control.Monad.State.Strict as State
import Data.Foldable (forM_)
import Control.Monad (when)

import Test.QuickCheck
import Test.QuickCheck.Monadic

Let’s define the abstract interface. v will be the type of variables that can operate on the monad m, holding values of type a. We need three operation: create a variable, read, and write to it.

class Var v m a | m -> v where
  new :: a -> m (v a)
  get :: v a -> m a
  set :: v a -> a -> m ()

Notice we had to use FunctionalDependencies to ease type inference. I don’t like this, there is probably a better way.

A utility function to do both read and write passing through a function:

modify :: (Monad m, Var v m a) => (a -> a) -> v a -> m ()
modify f v = get v >>= set v . f

Now we can provide different implementations for variables. First one in IO

newtype IOVar a = IOVar (IOArray () a)

We represent the value as an array of a single element. This is obviously overkill, but the goal was also to experiment with the low level array API

instance Var IOVar IO a where
  new = fmap IOVar . newArray ((), ())
  get (IOVar ar) = readArray ar ()
  set (IOVar ar) a = writeArray ar () a

The implementation is straightforward, reading and writing from/to the the array. Creation needs to take care of wrapping the array in the IOVar constructor. Notice that () is a valid index type for arrays, and it makes obvious in the type the fact that the array has a single element.

Providing an implementation in the ST monad is not much harder. Here, we could also use an STArray, but we go directly to STRef for simplicity

newtype STVar s a = STVar (STRef s a)

instance Var (STVar s) (ST s) a where
  new = fmap STVar . newSTRef
  get (STVar ref) = readSTRef ref
  set (STVar ref) a = writeSTRef ref a

The code looks very similar to the IO case.

Finally, let’s try to implement a variable in the State monad. For a variable holding values of type a, it is enough to maintain state a. So we can define

newtype StateVar a = StateVar a

And now to create an instance of Var we can do

instance Var StateVar (State.State a) a where
  new x = State.StateT $ \_ -> return (StateVar x, x)
  get _ = State.get
  set _ = State.put

get and set are simple. new requires some care. Initializing the variable means setting the state to a given value, so it can then be read by get. So in new we need it ignore the current state, and set it to x. The types are not enough to ensure correctness, there is a wrong implementation that also compiles:

wrongNew x = return . StateVar

And that’s it, we have the three types of variables we wanted. Now we can write a stateful looking algorithm, computing the maximum of a list is a good example. The way people do this in non functional languages usually is:

• initialize a variable max with the first element of the list
• go through all other elements:
  • if the current element is larger than max, update max with the new value
• when done iterating the list return max

We can express exactly this algorithm with our variables, even more, we can do it in a way that is generic for every type of Var and every supported Monad

myMaximum :: (Monad m, Var v m a, Ord a) => NonEmpty a -> m a

Take a look at the signature: given a non empty list (NonEmpty a), we return its maximum in some monad (Monad m). We can do this as long as a can be ordered (Ord a), and there is some type of variable v which works for the monad m and the type a (Var v m a). The type signature expresses all this pretty well.

myMaximum xs = do
  max <- new (NE.head xs) -- initialize a new var
  forM_ (NE.tail xs) $ \a -> do  -- for each el after the head
    maxSoFar <- get max  -- get the current maximum
    when (a > maxSoFar) $  -- compare with current element
      set max a  -- update if needed
  get max

Just like in the description of the algorithm, we create a variable and update it for every element that is larger than the initial value. When done iterating we return the last value hold by the variable.

Now we need to write some tests:

In IO

testIO :: (NonEmptyList Int) -> Property
testIO (NonEmpty xs) = monadicIO $ do
  mine <- run . myMaximum . NE.fromList $ xs
  assert $ mine == maximum xs

In ST

testST :: (NonEmptyList Int) -> Property
testST (NonEmpty xs) = monadicST $ do
  mine <- run . myMaximum . NE.fromList $ xs
  assert $ mine == maximum xs

And in State

testState :: (NonEmptyList Int) -> Bool
testState (NonEmpty xs) =
  State.execState (mine xs) whoCares == maximum xs
  where
    mine = myMaximum . NE.fromList
    whoCares = 42

Running the QuickCheck tests

main = do
  quickCheckWith opts testIO
  quickCheckWith opts testST
  quickCheckWith opts testState
  where opts = stdArgs {maxSuccess = 5000}

Success!

+++ OK, passed 5000 tests.
+++ OK, passed 5000 tests.
+++ OK, passed 5000 tests.