Can a `ST`-like monad be executed purely (without the `ST` library)?

Question

This post is literate Haskell. Just put in a file like \"pad.lhs\" and ghci will be able to run it.

> {-# LANGUAGE GADTs, Rank2Types #-}
>         


        

Answer 1

A simple solution is to wrap a State monad and present the same API as ST. In this case there's no need to store runtime type information, since it can be determined from the type of STRef-s, and the usual ST s quantification trick lets us prevent users from messing up the container storing the references.

We keep ref-s in an IntMap and increment a counter each time we allocate a new ref. Reading and writing just modifies the IntMap with some unsafeCoerce sprinkled atop.

{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving, RankNTypes, RoleAnnotations #-}

module PureST (ST, STRef, newSTRef, readSTRef, modifySTRef, runST) where

import Data.IntMap (IntMap, (!))
import qualified Data.IntMap as M

import Control.Monad
import Control.Applicative
import Control.Monad.Trans.State
import GHC.Prim (Any)
import Unsafe.Coerce (unsafeCoerce)

type role ST nominal representational
type role STRef nominal representational
newtype ST s a = ST (State (IntMap Any, Int) a) deriving (Functor, Applicative, Monad)
newtype STRef s a = STRef Int deriving Show

newSTRef :: a -> ST s (STRef s a)
newSTRef a = ST $ do
  (m, i) <- get
  put (M.insert i (unsafeCoerce a) m, i + 1)
  pure (STRef i)

readSTRef :: STRef s a -> ST s a
readSTRef (STRef i) = ST $ do
  (m, _) <- get
  pure (unsafeCoerce (m ! i))

writeSTRef :: STRef s a -> a -> ST s ()
writeSTRef (STRef i) a = ST $ 
  modify $ \(m, i') -> (M.insert i (unsafeCoerce a) m, i')

modifySTRef :: STRef s a -> (a -> a) -> ST s ()
modifySTRef (STRef i) f = ST $
  modify $ \(m, i') -> (M.adjust (unsafeCoerce f) i m, i')                      

runST :: (forall s. ST s a) -> a
runST (ST s) = evalState s (M.empty, 0)

foo :: Num a => ST s (a, Bool)
foo = do
  a <- newSTRef 0 
  modifySTRef a (+100)
  b <- newSTRef False
  modifySTRef b not
  (,) <$> readSTRef a <*> readSTRef b

Now we can do:

> runST foo
(100, True)

But the following fails with the usual ST type error:

> runST (newSTRef True)

Of course, the above scheme never garbage collects references, instead it frees up everything on each runST call. I think a more complex system could implement multiple distinct regions, each tagged by a type parameter, and allocate/free resources in a more fine-grained manner.

Also, the use of unsafeCoerce means here that using internals directly is every bit as dangerous as using GHC.ST internals and State# directly, so we should make sure to present a safe API, and also test our internals thoroughly (or else we may get segfaults in Haskell, a great sin).