Hi,

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DeriveFunctor #-}
-- Copyright 2013 Kevin Backhouse.
-- Copyright 2017 Timotej Tomandl. Modifications

import Data.STRef
import Control.Monad.ST
import Control.Applicative
import Control.Monad

newtype ST2 r w s a = ST2 { unwrapST2 :: ST s a }
    deriving (Functor)
instance Monad (ST2 r w s) where
  return x = ST2 $ return x
  (ST2 m) >>= f = ST2 $ do
        x <- m
        unwrapST2 (f x)
instance Applicative (ST2 r w s) where
  pure = return
  (<*>) = ap

-- | This function checks that the sub-computation is polymorphic in
-- both type parameters. This means that the sub-computation does not
-- read or write any state from the enclosing context.
{-# INLINE pureST2 #-}
pureST2 :: (forall r w. ST2 r w s a) -> ST2 r' w' s a
pureST2 m = m

-- | This function checks that the computation is polymorphic in
-- both parameters and then returns a pure value
{-# INLINE runST2 #-}
runST2 :: (forall r w s. ST2 r w s a) -> a
runST2 m=runST $ unwrapST2 m

-- | This function checks that the sub-computation is polymorphic in
-- the @w@ type parameter. This means that the sub-computation does
-- not write any state from the enclosing context (but read-only
-- operations are permitted).
{-# inline readOnlyST2 #-}
readOnlyST2 :: (forall w. ST2 r w s a) -> ST2 r w' s a
readOnlyST2 m = m

-- | Mutable reference. 'ST2Ref' is actually just a newtype of an
-- 'STRef', but the @r@ and @w@ type parameters allow the read and
-- write dependencies to be tracked by the type system.
newtype ST2Ref r w s a = ST2Ref (STRef s a)

-- | Create a new reference. The @r@ and @w@ type parameters of the
-- reference are unified with the 'ST2' monad to indicate that new
-- state is created in the enclosing context.
{-# INLINE newST2Ref #-}
newST2Ref :: a -> ST2 r w s (ST2Ref r w s a)
newST2Ref x = ST2 $ do
    r <- newSTRef x >>= \ var -> return var
    return (ST2Ref r)

-- | Read a reference. The @w@ type parameter of the reference is not
-- unified with the 'ST2' monad to indicate that this access is
-- read-only.
{-# INLINE readST2Ref #-}
readST2Ref :: ST2Ref r w s a -> ST2 r w' s a
readST2Ref (ST2Ref r) = ST2 $ readSTRef r

-- | Write to a reference. The @w@ type parameter of the reference is
-- unified with the 'ST2' monad to indicate that state is written in
-- the enclosing context.
{-# INLINE writeST2Ref #-}
writeST2Ref :: ST2Ref r w s a -> a -> ST2 r w s ()
writeST2Ref (ST2Ref r) x = ST2 $ writeSTRef r x

-- | Modify a reference.
{-# INLINE modifyST2Ref #-}
modifyST2Ref :: ST2Ref r w s a -> (a -> a) -> ST2 r w s ()
modifyST2Ref (ST2Ref r) f = ST2 $ modifySTRef r f