New monads/MonadSTO

From New monads, copied from old wiki.

Here's an extension of the ST monad in which the references are ordered and showable (they list their creation index).

{-# OPTIONS_GHC -fglasgow-exts #-}
module STO (
        STO,
        runSTO,
        newSTORef,
        readSTORef,
        writeSTORef,
        modifySTORef
    ) where

import Control.Monad.State
import Control.Monad.ST
import Data.STRef
import Data.Ord (comparing)

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 = comparing idx

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