FarTooGeneral

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

{- 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)

{- INSTRUCTION OPERATIONS -}

single :: (Eq i, Eq d) => Ins i d -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))
single PutState e = do
    (stack, map) <- tpop e
    stack' <- push e stack
    return (Branch stack' map)
single TakeState e = top e >>= (\state -> top state)
single Define e = do
    (stack, map) <- tpop e
    k <- (top >=> unwrapNode) stack
    (v, stack') <- (pop >=> tpop) stack
    map' <- appendMap k v map
    return (Branch stack' map')
single Lookup e = do
    (stack, map) <- tpop e
    k <- (top >=> unwrapNode) stack
    v <- indexMap k map
    stack' <- (pop >=> push v) stack
    return (Branch stack' map)
single Exec e = do
    (stack, map) <- tpop e
    (v, stack') <- tpop stack
    multiple v (Branch stack' map)
single d e = do
    stack <- top e
    map <- pop e
    stack' <- push (Node d) stack
    return (Branch stack' map)

multiple :: (Eq i, Eq d) => Elem (Ins i d) -> Elem (Ins i d) -> Either RunErr (Elem (Ins i d))
multiple (Branch (Node ins) next) state = single ins state >>= (\state' -> multiple next state')
multiple Leaf state = Right state -- return once program is finished
multiple _ _ = Left TypeErr