StateOnST.hs

{-# LANGUAGE 
    BangPatterns, 
    GADTs, 
    TypeFamilies, 
    FlexibleInstances, 
    UndecidableInstances, 
    MultiParamTypeClasses, 
    StandaloneDeriving, 
    DeriveFunctor, 
    DerivingVia, 
    RankNTypes, 
    ScopedTypeVariables 
    #-}

import Control.Monad.ST
import Data.STRef
import Control.Monad.Trans
import Control.Monad.State.Class
import Control.Monad.Reader

-- quoted from Ed Kmett's monad-st library
class Monad m => MonadST m where
  type World m :: *
  liftST :: ST (World m) a -> m a

instance MonadST IO where
  type World IO = RealWorld
  liftST = stToIO

instance MonadST (ST s) where
  type World (ST s) = s
  liftST = id

instance (MonadTrans t, MonadST m, Monad (t m)) => MonadST (t m) where
  type World (t m) = World m
  liftST = lift . liftST
-- end of quote

-- effectful MonadState transformer working on top of both ST and IO
newtype StateT s m a = StateT { runStateT :: STRef (World m) s -> m a }

deriving via ReaderT (STRef (World m) s) m instance Functor m =>
  Functor (StateT s m)

deriving via ReaderT (STRef (World m) s) m instance Applicative m =>
  Applicative (StateT s m)

deriving via ReaderT (STRef (World m) s) m instance Monad m =>
  Monad (StateT s m)

instance MonadST m => MonadState s (StateT s m) where
  get   = StateT $ \ref -> liftST $ readSTRef ref
  put v = StateT $ \ref -> liftST $ writeSTRef ref v

-- pure implementation for ST
runSTState :: s -> (forall w. StateT s (ST w) a) -> a
runSTState val act = runST $ run' val act
  where
    run' :: s -> (forall w. StateT s (ST w) a) -> (forall w. ST w a)
    run' val act = newSTRef val >>= runStateT act

-- impure implementation for IO
runIOState :: s -> StateT s (ST RealWorld) a -> IO a
runIOState val (StateT act) = stToIO $ do
  stref <- newSTRef val
  act stref

main = print $ runSTState 5 $ do
  val <- get
  put (val * 2)
  val <- get
  pure val