New monads/MonadSTO

{-# OPTIONS_GHC -fglasgow-exts #-}
module STO (
        STO,
        runSTO,
        newSTORef,
        readSTORef,
        writeSTORef,
        modifySTORef
    ) where
 
import Control.Monad.State
import Control.Monad.ST
import Data.STRef
 
newtype STO s a = STO { unSTO :: StateT Integer (ST s) a }
    deriving (Functor, Monad)
 
runSTO :: (forall s. STO s a) -> a
runSTO x = runST (evalStateT (unSTO x) 0)
 
data STORef s a = STORef { idx :: Integer, ref :: STRef s a }
    deriving Eq
 
instance Show (STORef s a) where
    show ref = "<STORef: " ++ (show (idx ref)) ++ ">"
 
instance Ord (STORef s a) where
    compare x y = compare (idx x) (idx y) 
 
newSTORef :: a -> STO s (STORef s a)
newSTORef x = STO $ do
    n <- get
    put (n+1)
    r <- lift $ newSTRef x
    return $ STORef { idx = n, ref = r }
 
readSTORef :: STORef s a -> STO s a
readSTORef (STORef { ref = r }) = STO $ lift $ readSTRef r
 
writeSTORef :: STORef s a -> a -> STO s ()
writeSTORef (STORef { ref = r }) x = STO $ lift $ writeSTRef r x
 
modifySTORef :: STORef s a -> (a -> a) -> STO s ()
modifySTORef (STORef { ref = r }) f = STO $ lift $ modifySTRef r f