Something like STRef but usable in a monad transformer

Counterexample.hs

{-# LANGUAGE RankNTypes #-}

-- Adapted from https://gist.github.com/Lysxia/18fce5d4ce8ff7eb7117dfb827770187 to work on StatesT

import Control.Monad
import Control.Monad.Trans.Class
import Data.Void
import StatesT

-- Demonstration that StatesT is unsafe

-- The continuation monad
newtype K r a = K ((a -> r) -> r)

instance Functor (K r) where
  fmap = liftM

instance Applicative (K r) where
  (<*>) = ap
  pure x = K (\k -> k x)

instance Monad (K r) where
  K u >>= h = K (\k -> u (\ x -> h' x k)) where
    h' x = let K y = h x in y

runStatesTK :: (forall s. StatesT s (K r) r) -> r
runStatesTK y =
  let K u = runStatesT y in u id

-- Sanity check that StatesT and K "work" in simple cases
okStates :: StatesT s (K r) (Int, Bool)
okStates = do
  x <- newStateRef (0 :: Int)
  y <- newStateRef True
  liftM2 (,) (readStateRef x) (readStateRef y)
  -- (0, True)

-- !!FUN!! EXPERIMENT BELOW

-- Use a time machine (e -> K r Void) in one timeline to send an e to
-- the other timeline.
--
-- How this works: first return a Right (starting the first timeline),
-- with a certain function (e -> _), which, if it is ever called
-- (with some argument e), backtracks to return the argument
-- as a Left (ending the first and starting the second timeline).
--
-- This is also known as a proof of the double negation of the
-- Law of Excluded Middle (LEM) in intuitionistic logic.
lem :: K r (Either e (e -> K r r'))
lem = K (\k -> k (Right (\e -> K (\_ -> k (Left e)))))

-- Breaking the type system with StatesT
badStates :: a -> StatesT s (K r) b
badStates a = do
  x <- lift lem  -- x :: Either (StateRef s b) (StateRef s b -> K r b)
  case x of
    Right k -> do
      -- In the first timeline, we are given a time machine k which can carry a (StateRef s b)
      r <- newStateRef undefined   -- We allocate a (StateRef s b)
      lift (k r)           -- We send it back in time
    Left r -> do
      -- In the second timeline, we get a reference out of nowhere (actually, from the previous timeline).
      _ <- newStateRef a     -- Allocate a reference of type a for **no reason** (actually, to populate the place the previous reference is pointing to with something of a different type)
      readStateRef r         -- Read the old reference

unsafeCoerce_ :: a -> b
unsafeCoerce_ a = runStatesTK (badStates a)

main :: IO ()
main = do
  print (runStatesTK okStates)   -- Sanity check (0, True)
  print (unsafeCoerce_ True :: Int)  -- segfault or a very bad looking number

Main.hs

import Control.Monad.Trans.Class (lift)
import StatesT

drawFrom :: StateRef s [a] -> StatesT s [] a
drawFrom urn = do
  urnContents <- readStateRef urn
  (chosenElement, remainingContents) <- lift $ removeEach urnContents
  writeStateRef urn remainingContents
  pure chosenElement

removeEach :: [a] -> [(a, [a])]
removeEach [] = []
removeEach (x:xs) = (x, xs):map (fmap (x:)) (removeEach xs)

result :: [(Int, Bool, Int)]
result = runStatesT $ do
  urn <- newStateRef [10,20,30,30]
  otherUrn <- newStateRef [True,False]
  x <- drawFrom urn
  b <- drawFrom otherUrn
  y <- drawFrom urn
  pure (x, b, y)

main :: IO ()
main = print result

-- [(10,True,20),(10,True,30),(10,True,30),(10,False,20),(10,False,30),(10,False,30),(20,True,10),(20,True,30),(20,True,30),(20,False,10),(20,False,30),(20,False,30),(30,True,10),(30,True,20),(30,True,30),(30,False,10),(30,False,20),(30,False,30),(30,True,10),(30,True,20),(30,True,30),(30,False,10),(30,False,20),(30,False,30)]

StatesT.hs

{-# LANGUAGE GeneralizedNewtypeDeriving, MagicHash, RankNTypes, RoleAnnotations #-}

module StatesT (StatesT, runStatesT, StateRef, newStateRef, readStateRef, writeStateRef, modifyStateRef, modifyStateRef') where

import Control.Applicative (Alternative)
import Control.Monad.State
import Data.List (genericIndex, genericSplitAt)
import GHC.Exts (Any, unsafeCoerce#)
import Numeric.Natural (Natural)

newtype StatesT s m a = StatesT (StateT (Natural, [Any]) m a) deriving(MonadTrans,Monad,Functor,Applicative,Alternative)
type role StatesT nominal nominal nominal -- I only want s to be nominal; m and a have to be nominal.

newtype StateRef s a = StateRef Natural
type role StateRef nominal representational

runStatesT :: Monad m => (forall s. StatesT s m a) -> m a
runStatesT (StatesT f) = evalStateT f (0, [])

newStateRef :: Monad m => a -> StatesT s m (StateRef s a)
newStateRef x = StatesT $ do
  (n, refs) <- get
  put (n + 1, refs ++ [unsafeCoerce# x])
  pure (StateRef n)

readStateRef :: Monad m => StateRef s a -> StatesT s m a
readStateRef (StateRef i) = StatesT $ unsafeCoerce# . (`genericIndex` i) . snd <$> get

writeStateRef :: Monad m => StateRef s a -> a -> StatesT s m ()
writeStateRef (StateRef i) x = StatesT . modify . fmap $ \refs -> let (before, _:after) = genericSplitAt i refs in before ++ unsafeCoerce# x : after

modifyStateRef :: Monad m => StateRef s a -> (a -> a) -> StatesT s m ()
modifyStateRef ref f = writeStateRef ref . f =<< readStateRef ref

modifyStateRef' :: Monad m => StateRef s a -> (a -> a) -> StatesT s m ()
modifyStateRef' ref f = do
  x <- readStateRef ref
  let x' = f x
  x' `seq` writeStateRef ref x'