Simulate ST Monad without "impure" things

with-intmap-dynamic.hs

#!/bin/env stack
{-
  stack script --resolver=lts-14.6
    --package=mtl
    --package=containers
-}

{-# OPTIONS -fdefer-type-errors #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE RankNTypes #-}

import           Data.Dynamic
import qualified Data.IntMap.Strict as IM
import           Data.Maybe
import           Control.Exception
import           Control.Monad.State.Strict

data STState = STState { lastId :: !Int, refs :: !(IM.IntMap Dynamic) } deriving Show

newtype STRef s a = STRef { unSTRef :: Int }

newtype ST s a = ST { unST :: State STState a }
  deriving (Functor, Applicative, Monad)

runST :: (forall s. ST s a) -> a
runST = (`evalState` (STState 0 IM.empty)) . unST

newSTRef :: Typeable a => a -> ST s (STRef s a)
newSTRef x = do
  state <- ST get
  let curId = lastId state
      curRefs = refs state
  ST $ put $! STState (curId + 1) $ IM.insert curId (toDyn x) curRefs
  return $ STRef curId

readSTRef :: Typeable a => STRef s a -> ST s a
readSTRef (STRef i) =
  ST $ gets
    ( fromMaybe (error $ "Assertion failure: Ref#" ++ show i ++ " is has unexpected type!")
    . fromDynamic
    . fromMaybe (error $ "Assertion failure: No ref#" ++ show i)
    . IM.lookup i
    . refs
    )

main :: IO ()
main = do
  print $ runST $ do
    ref1 <- newSTRef 'a'
    ref2 <- newSTRef True
    (,) <$> readSTRef ref1 <*> readSTRef ref2
  evaluate $ runST $ newSTRef True -- This line should cause a type error