{-# LANGUAGE GADTs #-}
module HW4Part2 where
{-
 - int ::= (any integer)
 -
 - reg ::= A | B | R
 -
 - bexpr ::= not bexpr
 -         | iexpr <= iexpr
 -
 - iexpr ::= int
 -         | reg
 -         | iexpr + iexpr
 -
 - stmt ::= reg ::= iexpr
 -        | if bexpr then prog else prog endprog else prog end
 -        | do prog end
 -        | break
 -
 - prog ::= \epsilon | stmt; prog
 -}

import Control.Monad.Except
import Control.Monad.State
import Data.Bool
import Data.Either

data Reg = A | B | R

data Expr a where
    Lit :: Int -> Expr Int
    Reg :: Reg -> Expr Int
    Plus :: Expr Int -> Expr Int -> Expr Int
    Leq :: Expr Int -> Expr Int -> Expr Bool
    Not :: Expr Bool -> Expr Bool

data Stmt
    = Store Reg (Expr Int)
    | If (Expr Bool) Prog Prog
    | Loop Prog
    | Break

type Prog = [Stmt]

-- Quick and dirty interpreter to test my programs :)

type Evaluator = ExceptT () (State (Int, Int, Int))

getReg :: Reg -> (Int, Int, Int) -> Int
getReg A (i, _, _) = i
getReg B (_, i, _) = i
getReg R (_, _, i) = i

setReg :: Reg -> Int -> (Int, Int, Int) -> (Int, Int, Int)
setReg A a (_, b, r) = (a, b, r)
setReg B b (a, _, r) = (a, b, r)
setReg R r (a, b, _) = (a, b, r)

expr :: Expr a -> Evaluator a
expr (Lit i) = return i
expr (Reg r) = gets (getReg r)
expr (Plus l r) = liftM2 (+) (expr l) (expr r)
expr (Leq l r) = liftM2 (<=) (expr l) (expr r)
expr (Not e) = not <$> expr e

stmt :: Stmt -> Evaluator ()
stmt (Store r e) = expr e >>= modify . setReg r
stmt (If i t e) = expr i >>= bool (prog e) (prog t)
stmt (Loop p) = prog (cycle p) `catchError` (const $ return ())
stmt Break = throwError ()

prog :: Prog -> Evaluator ()
prog = mapM_ stmt

run p = snd $ runState (runExceptT (prog p)) (0, 0, 0)