Rofth-FTG

import Control.Monad ((>=>))

data Elem a = Leaf | Node a | Branch (Elem a) (Elem a) deriving(Show)

data Ins i d
    = PutState
    | TakeState
    | Exec
    | Define
    | Lookup
    | Lit d
    | Ins i
    deriving(Eq, Show)

data RunErr
    = TypeErr
    | OutOfRange
    | InsErr
    deriving(Show)

toList :: [a] -> Elem a
toList [] = Leaf
toList (x:xs) = Branch (Node x) (toList xs)

unwrapNode :: Elem a -> Either RunErr a
unwrapNode (Node a) = Right a
unwrapNode _ = Left TypeErr

{- STACK/LIST OPERATIONS -}

top :: Elem a -> Either RunErr (Elem a)
top (Branch x _) = Right x
top Leaf = Left OutOfRange
top (Node _) = Left TypeErr

pop :: Elem a -> Either RunErr (Elem a)
pop (Branch _ xs) = Right xs
pop Leaf = Left OutOfRange
pop (Node _) = Left TypeErr

tpop :: Elem a -> Either RunErr (Elem a, Elem a) -- (top, pop)
tpop (Branch x xs) = Right (x, xs)
tpop Leaf = Left OutOfRange
tpop (Node _) = Left TypeErr

push :: Elem a -> Elem a -> Either RunErr (Elem a)
push x = Right . Branch x

append :: Elem a -> Elem a -> Either RunErr (Elem a)
append e (Branch x xs) = (append e xs) >>= (\xs' -> return (Branch x xs'))
append e Leaf = Right (Branch e Leaf)
append _ _ = Left TypeErr

updateFirst :: (Elem a -> Either RunErr (Elem a)) -> Elem a -> Either RunErr (Elem a)
updateFirst fn (Branch a b) = do
    v <- fn a
    return (Branch v b)
updateFirst _ _ = Left TypeErr

updateSecond :: (Elem a -> Either RunErr (Elem a)) -> Elem a -> Either RunErr (Elem a)
updateSecond fn (Branch a b) = fn b >>= (return . Branch a)
updateSecond _ _ = Left TypeErr

{- MAP OPERATIONS -}

indexMap :: (Eq a) => a -> Elem a -> Either RunErr (Elem a)
indexMap a map = do
    e <- top map
    case e of
        (Branch (Node a') e) | a == a' -> return e
        (Branch (Node _) _) -> do
            map' <- pop map
            indexMap a map'
        _ -> Left TypeErr

appendMap :: a -> Elem a -> Elem a -> Either RunErr (Elem a)
appendMap k v = append (Branch (Node k) v)

{- STATE OPERATIONS -}
incIns :: Elem (Ins i d) -> Either RunErr (Elem (Ins i d)) -- increment ins
incIns s = updateFirst (updateFirst pop) s

retProc :: Elem (Ins i d) -> Either RunErr (Elem (Ins i d)) -- return from process
retProc s = updateFirst pop s

nextIns :: Elem (Ins i d) -> Either RunErr (Ins i d, Elem (Ins i d)) -- next ins, new state
nextIns s = case (top >=> top >=> top >=> unwrapNode) s of
    Right ins -> incIns s >>= (return . (,) ins)
    Left err -> retProc s >>= nextIns

single :: (Eq i, Eq d) => (i -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))) -> (Ins i d) -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))
single fn (Ins i) s = fn i s
single _ PutState s = updateSecond (updateFirst (push s)) s
single _ TakeState s = (pop >=> top >=> top) s
single _ Exec s = do
    proc <- (pop >=> top >=> top) s
    updateFirst (push proc) s
single _ Define s = do
    k <- (pop >=> top >=> top >=> unwrapNode) s
    v <- (pop >=> top >=> pop >=> top) s
    updateSecond
        (   updateFirst (pop >=> pop)
        >=> updateSecond (appendMap k v)) s
single _ Lookup s = do
    k <- (pop >=> top >=> top >=> unwrapNode) s
    v <- (pop >=> pop >=> indexMap k) s
    updateSecond
        (updateFirst (pop >=> push v)) s
single _ i@(Lit _) s = updateSecond (updateFirst (push (Node i))) s

doStep :: (Eq i, Eq d) => (i -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))) -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))
doStep fn s = do
    (ins, s') <- nextIns s
    single fn ins s'

runPrg :: (Eq i, Eq d) => (i -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))) -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))
runPrg fn s = case top s of
    Right (Branch Leaf Leaf) -> return s
    _ -> do
        s' <- doStep fn s
        runPrg fn s'

load :: [Ins i d] -> Elem (Ins i d)
load ls = Branch (Branch (toList ls) Leaf) (Branch Leaf Leaf)

showOutput :: (Show i, Show d) => Elem (Ins i d) -> String
showOutput = \y -> case (pop >=> top >=> top) y of
    Right x -> show x
    _ -> "No top item on stack"

execPrg :: (Eq i, Eq d, Show i, Show d) => [Ins i d] -> String -- now you can do 'execPrg [Lit "hello, world", ...]' etc.
execPrg ls = case runPrg undefined (load ls) of
    Right out -> showOutput out
    Left err -> show err