A version of http://d.hatena.ne.jp/fumiexcel/20121111/1352624678 that protects against labels escaping an invocation of runGotoT

Goto.hs

{-# LANGUAGE RankNTypes #-}
import qualified Data.IntMap as M
import Control.Monad.Trans.Free
import Control.Monad.Trans

{-
The phantom parameter s protects labels from escaping an invocation of runGotoT.
For example, this program:
invalid = runGotoT label >>= runGotoT . goto
is rejected at compile time with this version, but crashes
at run time in the original version at http://d.hatena.ne.jp/fumiexcel/20121111/1352624678.

This trick is inspired by the ST monad in base: https://hackage.haskell.org/package/base-4.11.0.0/docs/Control-Monad-ST.html

-}

newtype Label s = LabelId Int deriving (Eq, Ord)

data Labeling s a = Label (Label s -> a) | Goto (Label s)

instance Functor (Labeling s) where
    fmap f (Label g) = Label (f . g)
    fmap f (Goto l) = Goto l

type GotoT s = FreeT (Labeling s)

label :: Monad m => GotoT s m (Label s)
label = liftF (Label id)

goto :: Monad m => Label s -> GotoT s m ()
goto l = liftF (Goto l)

runGotoT :: Monad m => (forall s. GotoT s m a) -> m a
runGotoT = run M.empty where
    run st m = runFreeT m >>= \r -> case r of
        Pure a -> return a
        Free (Label f) -> let cont = f (LabelId newLabel) in run (M.insert newLabel cont st) cont
        Free (Goto (LabelId i)) -> run st (st M.! i)
        where
            newLabel = succ (M.size st)

ex2 = do
    lift $ putStrLn "Begin."
    
    hoge <- label
    lift $ putStrLn "Label hoge."
    
    piyo <- label
    lift $ putStrLn "Label piyo."
    
    fuga <- label
    lift $ putStrLn "Label fuga."
        
    lift $ putStr "Where do you want to go? "
    str <- lift getLine
    case str of
        "hoge" -> goto hoge
        "piyo" -> goto piyo
        "fuga" -> goto fuga
        _ -> lift $ putStrLn "Quit."

main = runGotoT ex2