Modeling the ST monad in Agda

This recent SO question prompted me to write an unsafe and pure emulation of the ST monad in Haskell, a slightly modified version of which you can see below:

{-# LANGUAGE DeriveFunctor, GeneralizedNewtypeDeriving, RankNTypes #-}

import Control.Monad.Trans.State
import GHC.Prim (Any)
import Unsafe.Coerce (unsafeCoerce)
import Data.List

newtype ST s a = ST (State ([Any], Int) a) deriving (Functor, Applicative, Monad)
newtype STRef s a = STRef Int deriving Show

newSTRef :: a -> ST s (STRef s a)
newSTRef a = ST $ do
  (env, i) <- get
  put (unsafeCoerce a : env, i + 1)
  pure (STRef i)

update :: [a] -> (a -> a) -> Int -> [a]
update as f i = case splitAt i as of
  (as, b:bs) -> as ++ f b : bs
  _          -> as

readSTRef :: STRef s a -> ST s a
readSTRef (STRef i) = ST $ do
  (m, i') <- get
  pure (unsafeCoerce (m !! (i' - i - 1)))

modifySTRef :: STRef s a -> (a -> a) -> ST s ()
modifySTRef (STRef i) f = ST $
  modify $ \(env, i') -> (update env (unsafeCoerce f) (i' - i - 1), i')

runST :: (forall s. ST s a) -> a
runST (ST s) = evalState s ([], 0)

It would be good if we could present the usual ST monad API without unsafeCoerce. Specifically, I'd like to formalize the reasons why the usual GHC ST monad and the above emulation works. It seems to me that they work because:

• Any STRef s a with the right s tag must have been created in the current ST computation, since runST ensures that different states can't be mixed up.
• The previous point together with the fact that ST computations can only ever extend the environment of references implies that any STRef s a with the right s tag refers to a valid a-typed location in the environment (after possibly weakening the reference at runtime).

The above points enable a remarkably proof-obligation-free programming experience. Nothing really comes close in safe and pure Haskell that I can think of; we can do a rather poor imitation with indexed state monads and heterogeneous lists, but that doesn't express any of the above points and thus requires proofs at each use site of STRef-s.

I'm at a loss how we could formalize this properly in Agda. For starters, "allocated in this computation" is tricky enough. I thought about representing STRef-s as proofs that a particular allocation is contained in a particular ST computation, but that seems to result in an endless recursion of type indexing.

Here's some form of a solution done by postulating a parametricity theorem. It also ensures that the postulate does not get in the way of computation.

http://code.haskell.org/~Saizan/ST/ST.agda

"darcs get http://code.haskell.org/~Saizan/ST/" for the full source

I'm not happy about the closed universe of types but it's the easy way to tailor the parametricity theorem to what we actually need.