Roman Cheplyaka

Universally stateless monads

June 4, 2017

Background: monad-control and stateless monads

The monad-control package allows to lift IO functions such as

forkIO :: IO () -> IO ThreadId

catch :: Exception e => IO a -> (e -> IO a) -> IO a

allocate :: MonadResource m => IO a -> (a -> IO ()) -> m (ReleaseKey, a)

to IO-based monad stacks such as StateT Int (ReaderT Bool IO).

The core idea of the package is the associated type StM, which, for a given monad m and result type a, calculates the “state” of m at a.

The “state” of a monad is whatever the “run” function for this monad returns.

For instance, for StateT Int IO Char, we have

runStateT :: StateT Int IO Char -> Int -> IO (Char, Int)

The result type of this function (minus the outer monad constructor, IO, which is always there) is (Char, Int), and that is what StM (StateT Int IO) Char should expand to:

> :kind! StM (StateT Int IO) Char
StM (StateT Int IO) Char :: *
= (Char, Int)

In this case, StM m a is not the same as a; it contains a plus some extra information.

In other cases, StM m a may not contain an a at all; for instance

> :kind! StM (ExceptT Text IO) Char
StM (ExceptT Text IO) Char :: *
= Either Text Char

and we cannot always extract a Char from Either Text Char.

For some monads, though, StM m a reduces precisely to a. I call such monads “stateless”. A notable example is the reader monad:

> :kind! StM (ReaderT Int IO) Bool
StM (ReaderT Int IO) Bool :: *
= Bool

Note that a monad like ReaderT (IORef Int) IO is also stateless, even though one can use it to implement stateful programs.

The important feature of stateless monads is that we can fork them without duplicating the state and terminate them without losing the state. The monad-control package works best with stateless monads: it is less tricky to understand, and you can do some things with stateless monads that are hard or impossible to do with arbitrary MonadBaseControl monads.

Universally stateless monads

When both the monad m and the result type a are known, the compiler can expand the associated synonym StM m a and decide whether StM m a ~ a.

However, there are good reasons to keep the monad m polymorphic and instead impose the constraints (e.g. MonadReader Int m) that m must satisfy.

In this case, the compiler cannot know a priori that m is stateless, and we need to explicitly state that in the function signature. In Taking advantage of type synonyms in monad-control, I showed one such example: running a web application with access to the database. In order to convince the compiler that m is stateless, I needed to add the constraint

StM m ResponseReceived ~ ResponseReceived

to the type signature.

As you can see, this doesn’t quite say “monad m is stateless”; instead it says “monad m is stateless at type a” (where a is ResponseReceived in the above example).

This is fine if we only use monad-control at one result type. But if we use monad-control functions at many different types, the number of constraints required quickly gets out of hand.

As an example, consider the allocate function from resourcet’s Control.Monad.Trans.Resource:

allocate :: MonadResource m => IO a -> (a -> IO ()) -> m (ReleaseKey, a)

As the module’s documentation says,

One point to note: all register cleanup actions live in the IO monad, not the main monad. This allows both more efficient code, and for monads to be transformed.

In practice, it is often useful for the register and cleanup actions to live in the main monad. monad-control lets us lift the allocate function:

{-# LANGUAGE FlexibleContexts, TypeFamilies, ScopedTypeVariables #-}
import Control.Monad.Trans.Control
import Control.Monad.Trans.Resource

allocateM
  :: forall m a . (MonadBaseControl IO m, MonadResource m,
                   StM m a ~ a, StM m () ~ (), StM m (ReleaseKey, a) ~ (ReleaseKey, a))
  => m a -> (a -> m ()) -> m (ReleaseKey, a)
allocateM acquire release =
  liftBaseWith
    (\runInIO -> runInIO $ allocate
      (runInIO acquire)
      (runInIO . release))

This small function requires three different stateless constraints — constraints of the form StM m x ~ x — and two additional stateless constraints per each additional type a at which we use it.

These constraints are artefacts of the way type synonyms work in Haskell; StM m a is not really supposed to depend on a. If a monad is stateless, it is stateless at every result type.

In order to capture this universal statelessness in a single constraint, we can use Forall from the constraints package.

First, we need to transform the constraint StM m a ~ a to a form in which it can be partially applied, as we want to abstract over a. Simply saying

type StatelessAt m a = StM m a ~ a

won’t do because type synonyms need to be fully applied: StatelessAt m is not a valid type.

We need to use a trick to create a class synonym:

class    StM m a ~ a => StatelessAt m a
instance StM m a ~ a => StatelessAt m a

Now StatelessAt is a legitimate class constraint (not a type synonym), and so we can abstract over its second argument with Forall:

type Stateless m = Forall (StatelessAt m)

Now we only need to include the Stateless m constraint, and every time we need to prove that StM m a ~ a for some result type a, we wrap the monadic computation in assertStateless @a (...), where assertStateless is defined as follows:

assertStateless :: forall a m b . Stateless m => (StM m a ~ a => m b) -> m b
assertStateless action = action \\ (inst :: Stateless m :- StatelessAt m a)

The type signature of assertStateless is crafted in such a way that we only need to specify a, and m is inferred to be the “current” monad. We could have given assertStateless a more general type

assertStateless :: forall m a r . Stateless m => (StM m a ~ a => r) -> r

but now we have to apply it to both m and a.

As an example of using assertStateless, let’s rewrite the lifted allocate function to include a single stateless constraint:

allocateM
  :: forall m a . (MonadBaseControl IO m, MonadResource m, Stateless m)
  => m a -> (a -> m ()) -> m (ReleaseKey, a)
allocateM acquire release =
  assertStateless @a $
  assertStateless @() $
  assertStateless @(ReleaseKey, a) $
  liftBaseWith
    (\runInIO -> runInIO $ allocate
      (runInIO acquire)
      (runInIO . release))

Here, assertStateless generated all three StM m x ~ x constraints for us on the fly, from the single universally-quantified constraint Stateless m.

Stateless monad transformers

Let’s say we are writing a function that works in some stateless monad, m:

foo :: (MonadBaseControl IO m, MonadResource m, Stateless m) => m ()
foo = do ...

But locally, it adds another layer on top of m:

foo :: (MonadBaseControl IO m, MonadResource m, Stateless m) => m ()
foo = do
  thing <- getThing
  flip runReaderT thing $ do
    ...

And somewhere in there we need to allocate something:

foo :: (MonadBaseControl IO m, MonadResource m, Stateless m) => m ()
foo = do
  thing <- getThing
  flip runReaderT thing $ do
    resource <- allocateM acq rel
    ...

The compiler won’t accept this, though:

• Could not deduce: StM
                      m (Data.Constraint.Forall.Skolem (StatelessAt (ReaderT () m)))
                    ~
                    Data.Constraint.Forall.Skolem (StatelessAt (ReaderT () m))
    arising from a use of ‘allocateM’
  from the context: (MonadBaseControl IO m,
                     MonadResource m,
                     Stateless m)
    bound by the type signature for:
               foo :: (MonadBaseControl IO m, MonadResource m, Stateless m) =>
                      m ()

In order to run allocateM in the inner environment, ReaderT Thing m, we need to satisfy the constraint Stateless (ReaderT Thing), which is different from the Stateless m that we have in scope.

If the acq and rel actions do not need to access the thing, we can avoid the problem by lifting the action to the outer environment:

foo :: (MonadBaseControl IO m, MonadResource m, Stateless m) => m ()
foo = do
  thing <- getThing
  flip runReaderT thing $ do
    resource <- lift $
      -- this happens in m
      allocateM acq rel
    ...

But what if acq and rel do need to know the thing?

In that case, we need to prove to the compiler that for all m, Stateless m implies Stateless (ReaderT Thing). This should follow from the fact that ReaderT e is itself a “stateless transformer”, meaning that it doesn’t change the state of the monad that it transforms. As with Stateless, we put this in the form of partially-applicable class and then abstract over a (and m):

class    StM (t m) a ~ StM m a => StatelessTAt t (m :: * -> *) a
instance StM (t m) a ~ StM m a => StatelessTAt t m a

type StatelessT t = ForallV (StatelessTAt t)

Now we need to prove that StatelessT t and Stateless m together imply Stateless (t m). In the notation of the constraints package, this statement can be written as

statelessT :: forall t m . (StatelessT t, Stateless m) :- Stateless (t m)

How to prove it in Haskell is not completely obvious, and I recommend that you try it yourself. I also posted a simplified version of this exercise on twitter the other day.

Anyway, here is one possible answer:

statelessT = Sub $ forall prf
  where
    prf :: forall a . (StatelessT t, Stateless m) => Dict (StatelessAt (t m) a)
    prf = Dict \\ (instV :: StatelessT t :- StatelessTAt t m a)
               \\ (inst  :: Stateless m  :- StatelessAt m a)

Finally, here is a function analogous to assertStateless that brings the Stateless (t m) constraint into scope:

liftStatelessT
  :: forall t m b . (StatelessT t, Stateless m)
  => (Stateless (t m) => (t m) b) -> (t m) b
liftStatelessT action = action \\ statelessT @t @m

And here is a minimal working example that demonstrates the usage of liftStatelessT:

foo :: (MonadBaseControl IO m, MonadResource m, Stateless m) => m ()
foo = do
  flip runReaderT () $ liftStatelessT $ do
    resource <- allocateM (return ()) (const $ return ())
    return ()

Complete code for Stateless and StatelessT

{-# LANGUAGE GADTs, ConstraintKinds, MultiParamTypeClasses, FlexibleInstances,
             ScopedTypeVariables, RankNTypes, AllowAmbiguousTypes,
             TypeApplications, TypeOperators, KindSignatures,
             UndecidableInstances, UndecidableSuperClasses #-}

module Stateless where

import Data.Constraint
import Data.Constraint.Forall
import Control.Monad.Trans.Control

class    StM m a ~ a => StatelessAt m a
instance StM m a ~ a => StatelessAt m a

-- | A constraint that asserts that a given monad is stateless
type Stateless m = Forall (StatelessAt m)

-- | Instantiate the stateless claim at a particular monad and type
assertStateless :: forall a m b . Stateless m => (StM m a ~ a => m b) -> m b
assertStateless action = action \\ (inst :: Stateless m :- StatelessAt m a)

class    StM (t m) a ~ StM m a => StatelessTAt t (m :: * -> *) a
instance StM (t m) a ~ StM m a => StatelessTAt t m a

-- | A statement that a monad transformer doesn't alter the state type
type StatelessT t = ForallV (StatelessTAt t)

-- | A proof that if @t@ is a stateless transformer and @m@ is a stateless monad,
-- then @t m@ is a stateless monad
statelessT :: forall t m . (StatelessT t, Stateless m) :- Stateless (t m)
statelessT = Sub $ forall prf
  where
    prf :: forall a . (StatelessT t, Stateless m) => Dict (StatelessAt (t m) a)
    prf = Dict \\ (instV :: StatelessT t :- StatelessTAt t m a)
               \\ (inst  :: Stateless m  :- StatelessAt m a)

-- | Derive the 'Stateless' constraint for a transformed monad @t m@
liftStatelessT
  :: forall t m b . (StatelessT t, Stateless m)
  => (Stateless (t m) => (t m) b) -> (t m) b
liftStatelessT action = action \\ statelessT @t @m