goto in Haskell

{-# LANGUAGE GADTs #-}

import Data.Map as Map
import Control.Monad.Trans.Either
import qualified Control.Monad.Trans.Class as Trans
import Control.Monad (liftM)

data GoTo k m r where
  Label :: k -> GoTo k m ()
  GoTo :: k -> GoTo k m ()
  Lift :: m a -> GoTo k m ()
  Return :: a -> GoTo k m a
  Seq :: GoTo k m () -> GoTo k m b -> GoTo k m b

showgo :: (Show k, Show r) => GoTo k m r -> String
showgo (Label k) = "  " ++ show k ++ " <---"
showgo (GoTo k)  = "  " ++ "---> " ++ show k
showgo (Lift m)  = "  prim"
showgo (Return a) = "  val " ++ show a
showgo (Seq m1 m2) = showgo m1 ++ ";\n" ++ showgo m2

instance (Show k, Show r) => Show (GoTo k m r) where
  show p = "Program {\n" ++ showgo p ++ ";\n}"

label :: k -> GoTo k m ()
label = Label

lift :: m a -> GoTo k m ()
lift = Lift

goto :: k -> GoTo k m ()
goto = GoTo

skip :: Monad m => m ()
skip = return ()

instance Monad m => Monad (GoTo k m) where
  return x = Return x
  m0 >>= f = case m0 of
    Label k   -> Seq (Label k) (f ())
    GoTo k    -> Seq (GoTo k) (f ())
    Lift m    -> Seq (Lift m) (f ())
    Return x  -> f x
    Seq m1 m2 -> Seq m1 (m2 >>= f)

labels :: (Ord k) => GoTo k m () -> Map k (GoTo k m ())
labels = go Map.empty (Return ()) where
  go ls rest (Label k) = Map.insert k (Seq (Label k) rest) ls
  go ls _ (GoTo _) = ls
  go ls _ (Lift _) = ls
  go ls _ (Return _) = ls
  go ls rest (Seq m1 m2) = go ls' rest m2 where
    ls' = go ls (Seq m2 rest) m1

data TerminationCode k
  = Success
  | NoSuchLabel k
  deriving Show

usesGoTo :: (Ord k, Monad m) => GoTo k m () -> m (TerminationCode k)
usesGoTo m0 = liftM fromEither $ runEitherT $ go m0 where
  go (Label _) = skip
  go (GoTo k) = case Map.lookup k ls of
    Nothing -> EitherT $ return $ Left k
    Just there -> go there -- Find it on the map? Go there. :)
  go (Lift m) = Trans.lift m >> skip
  go (Return ()) = skip
  go (Seq m1 m2) = case m1 of
    GoTo{} -> go m1 -- goto discards its 'current continuation' m2
    _ -> go m1 >> go m2

  -- static analysis!
  ls = labels m0

  fromEither (Left k) = NoSuchLabel k
  fromEither (Right ()) = Success

example = do
  goto (1 :: Int)
  label 2
  lift $ putStrLn "two"
  goto 3
  label 1
  lift $ putStrLn "one"
  goto 2
  label 3
  lift $ putStrLn "three"

main = usesGoTo example