module CS4400.State where

import Control.Monad (ap, liftM)

data Expr
  = ENum Int
  | EAdd Expr Expr
  deriving (Show)

type Addr = Int
type Reg = Int
type Value = Int

data Inst
  = ISet Reg Value
  | IStore Reg Addr
  | ILoad Reg Addr
  | IAdd
  | IHalt
  deriving (Show)

newtype State s a = State {runState :: s -> (s, a)}

instance Monad (State s) where
  (>>=) = andThen

instance Applicative (State s) where
  pure = yield
  (<*>) = ap

instance Functor (State s) where
  fmap = liftM

andThen :: State s a -> (a -> State s b) -> State s b
andThen s f =
  State
    ( \c ->
        let (c', x) = runState s c
         in runState (f x) c'
    )

get :: State s s
get = State (\c -> (c, c))

put :: s -> State s ()
put c = State (\_ -> (c, ()))

yield :: a -> State s a
yield x = State (\c -> (c, x))

fresh :: State Addr Addr
fresh = do
  counter <- get
  put (counter + 1)
  yield counter

compile :: Expr -> State Addr ([Inst], Addr)
compile expr = case expr of
  ENum n -> do
    addr <- fresh
    let insts = [ISet 0 n, IStore 0 addr]
    yield (insts, addr)
  EAdd e1 e2 -> do
    (insts1, addr1) <- compile e1
    (insts2, addr2) <- compile e2
    addr <- fresh
    let insts = [ILoad 1 addr1, ILoad 2 addr2, IAdd, IStore 0 addr]
    yield (insts1 ++ insts2 ++ insts, addr)

compileHalt :: Expr -> [Inst]
compileHalt expr =
  let (_, (insts, addr)) = runState (compile expr) 0
   in (insts ++ [ILoad 0 addr, IHalt])