Haskell monadic COMEFROM

comefrom.hs

{-# LANGUAGE RankNTypes, PatternSynonyms, GADTs, ViewPatterns, LambdaCase, ScopedTypeVariables #-}
module Main where
    
import Data.Map as Map
import Control.Monad (ap, liftM)
import Data.IORef

data ComeFromOp a where
    Line :: Int -> ComeFromOp ()
    LiftIO :: IO () -> ComeFromOp ()
    ComeFrom :: Int -> ComeFromOp ()


data ComeFrom a where
    Embed :: ComeFromOp a -> ComeFrom a
    Pure :: a -> ComeFrom a
    Bind :: ComeFrom a -> (a -> ComeFrom b) -> ComeFrom b

line :: Int -> ComeFrom ()
line = Embed . Line

liftIO :: IO () -> ComeFrom ()
liftIO = Embed . LiftIO

pattern COMEFROM :: Int -> ComeFrom ()
pattern COMEFROM i <- (error "This is only a pattern synonym to get the sick uppercase name" -> i)
    where
        COMEFROM i = Embed (ComeFrom i)



instance Functor ComeFrom where
    fmap = liftM
instance Applicative ComeFrom where
    (<*>) = ap
    pure = Pure
instance Monad ComeFrom where
    return = pure
    (>>=) = Bind


runComeFrom :: ComeFrom () -> IO ()
runComeFrom comefrom = do
    let comefromMap = collectComefroms comefrom
    go comefromMap comefrom
        where
            collectComefroms :: ComeFrom () -> Map Int (ComeFrom ())
            collectComefroms = \case
                Bind (Embed (ComeFrom i)) cont -> fromList [(i, cont ())] <> collectComeFroms (cont ()) 
                (Embed _) -> mempty
                (Pure _) -> mempty
                Bind (Embed (Line _)) cont -> collectComefroms (cont ())
                Bind (Embed (LiftIO _)) cont -> collectComefroms (cont ())
                Bind (Pure a) cont -> collectComefroms (cont a)
                Bind (Bind _ _) _cont -> error "nested binds are not supported"

            go :: Map Int (ComeFrom ()) -> ComeFrom () -> IO ()
            go comeFroms = \case 
                Pure x -> pure x
                Embed (ComeFrom _) -> pure ()
                Embed (LiftIO io) -> io
                Embed (Line line) -> pure ()
                Bind (Embed (Line line)) cont -> do
                    case Map.lookup line comeFroms of
                        Nothing -> go comeFroms (cont ())
                        Just lineCont ->
                            go comeFroms lineCont
                Bind (Embed (LiftIO io)) cont -> do
                    io
                    go comeFroms $ cont ()
                Bind (Embed (ComeFrom _)) cont -> do
                    go comeFroms $ cont ()
                Bind (Pure x) cont -> do
                    go comeFroms $ cont x
                Bind (Bind _ _) _cont -> error "nested binds are not supported"


main :: IO ()
main = runComeFrom $ do
    line 1
    liftIO $ putStrLn "aaa"
    line 2
    liftIO $ putStrLn "bbb"
    line 3
    COMEFROM 2
    line 4
    liftIO $ putStrLn "ccc"
    
{-
Result:

aaa
ccc

-}