data Var = X | Y | Z deriving (Eq, Show) 

{-
type State = Var -> Integer

lkp :: Var -> State -> Integer
lkp x s = s x

update :: Var -> Integer -> State -> State
update x z s = \ y -> if x == y then z else s y

init :: State
init = \ _ -> 0
-}

type State = [(Var, Integer)]

lkp :: Var -> State -> Integer
lkp x [] = 0
lkp x ((y,z):s) = if x == y then z else lkp x s

upd :: Var -> Integer -> State -> State
--upd x z s = (x,z):s 
upd x z [] = [(x,z)]
upd x z ((y,w):s) = if x == y then (x,z):s
                       else (y,w):upd x z s

init :: State
init = []


data AExp = N Integer | V Var 
          | AExp :+ AExp | AExp :- AExp | AExp :* AExp deriving Show

aexp :: AExp -> State -> Integer
aexp (N z) _      = z
aexp (V x) s      = lkp x s
aexp (a0 :+ a1) s = aexp a0 s + aexp a1 s
aexp (a0 :- a1) s = aexp a0 s - aexp a1 s
aexp (a0 :* a1) s = aexp a0 s * aexp a1 s


data BExp = TT | FF | AExp :== AExp | AExp :<= AExp 
          | Not BExp | BExp :&& BExp | BExp :|| BExp deriving Show

bexp :: BExp -> State -> Bool
bexp TT _ = True
bexp FF _ = False
bexp (a0 :== a1) s = aexp a0 s == aexp a1 s
bexp (a0 :<= a1) s = aexp a0 s <= aexp a1 s
bexp (Not b) s = not (bexp b s)
bexp (a0 :&& a1) s = bexp a0 s && bexp a1 s
bexp (a0 :|| a1) s = bexp a0 s || bexp a1 s


data Stmt = Skip | Stmt :\ Stmt 
          | Var := AExp
          | If BExp Stmt Stmt
          | While BExp Stmt deriving Show

stmt :: Stmt -> State -> State
stmt Skip s = s
stmt (stm0 :\ stm1) s = s' where
                          s'' = stmt stm0 s           
                          s'  = stmt stm1 s'' 
-- stmt (stm0 :\ stm1) s = stmt stm1 (stmt stm0 s)
stmt (x := a) s = upd x z s where
                          z = aexp a s
stmt (If b stm0 stm1) s = 
    if bexp b s then 
      stmt stm0 s
    else stmt stm1 s
stmt (While b stm0) s = 
    if bexp b s then 
      stmt (While b stm0) (stmt stm0 s) 
    else s

fac :: Stmt 
fac = (Y := N 1) :\ 
      (While (N 1 :<= V X)
        ((Y := (V Y :* V X)) :\
         (X := (V X :- N 1))  
        )
      )

-- abstract machines

type Config = (Label, Stack, State)

type Stack = [Either Integer Bool]

type Label = Integer

data Instr = Push Integer | Add | Sub | Mul
           | PushTrue | PushFalse | Eq | Le | Neg | And | Or
           | Fetch Var | Store Var
           | Jump Label | JumpFalse Label 
                                          deriving Show

type Code = [(Label, Instr)]


compAExp :: AExp -> Label -> (Code, Label)
compAExp (N z) ell = ([(ell, Push  z)], ell+1)
compAExp (V x) ell = ([(ell, Fetch x)], ell+1)
compAExp (a0 :+ a1) ell = 
     let (c0, ell' ) = compAExp a1 ell 
         (c1, ell'') = compAExp a0 ell' 
     in (c0 ++ c1 ++ [(ell'', Add)], ell''+1)
compAExp (a0 :- a1) ell = 
     let (c0, ell' ) = compAExp a1 ell 
         (c1, ell'') = compAExp a0 ell' 
     in (c0 ++ c1 ++ [(ell'', Sub)], ell''+1)
compAExp (a0 :* a1) ell = 
     let (c0, ell' ) = compAExp a1 ell 
         (c1, ell'') = compAExp a0 ell' 
     in (c0 ++ c1 ++ [(ell'', Mul)], ell''+1)

compBExp :: BExp -> Label -> (Code, Label)
compBExp TT ell = ([(ell, PushTrue )], ell+1)
compBExp FF ell = ([(ell, PushFalse)], ell+1)
compBExp (a0 :== a1) ell = 
     let (c0, ell' ) = compAExp a1 ell 
         (c1, ell'') = compAExp a0 ell' 
     in (c0 ++ c1 ++ [(ell'', Eq)], ell''+1)
compBExp (a0 :<= a1) ell = 
     let (c0, ell' ) = compAExp a1 ell 
         (c1, ell'') = compAExp a0 ell' 
     in (c0 ++ c1 ++ [(ell'', Le)], ell''+1) 
--compBExp (Not b) = compBExp b ++ [Neg]
--compBExp (b0 :&& b1) = compBExp b1 ++ compBExp b0 ++ [And]
--compBExp (b0 :|| b1) = compBExp b1 ++ compBExp b0 ++ [Or]


compStmt :: Stmt -> Label -> (Code, Label)
compStmt Skip ell = ([], ell)
compStmt (stm0 :\ stm1) ell =
     let (c0, ell' ) = compStmt stm0 ell
         (c1, ell'') = compStmt stm1 ell'
     in (c0 ++ c1, ell'')
compStmt (x := a) ell = 
     let (c, ell') = compAExp a ell
     in (c ++ [(ell', Store x)], ell'+1)
compStmt (If b stm0 stm1) ell
   = let (c , ell'  ) = compBExp b    ell
         (c0, ell'' ) = compStmt stm0 (ell' +1)
         (c1, ell''') = compStmt stm1 (ell''+1)
     in (c ++ [(ell', JumpFalse (ell''+1))] ++ c0 
           ++ [(ell'', Jump ell''')] ++ c1, ell''')
compStmt (While b stm0) ell
   = let (c , ell'  ) = compBExp b    ell
         (c0, ell'' ) = compStmt stm0 (ell' +1)
     in (c ++ [(ell', JumpFalse (ell''+1))] ++ c0
           ++ [(ell'', Jump ell)], ell''+1) 


makeStep :: Code -> Config -> Maybe Config 
makeStep c cfg@(ell, e, s) =
        case fetch ell c of
           Just instr -> Just (exec instr cfg)
           Nothing    -> Nothing 

fetch :: Label -> Code -> Maybe Instr
fetch _ [] = Nothing
fetch ell ((m, instr):c) = if ell == m then Just instr 
                           else fetch ell c

exec :: Instr -> Config -> Config
exec (Push z) (ell, e, s) = (ell+1, Left z : e, s)
exec Add (ell, Left z0 : Left z1 : e, s)
                            = (ell+1, Left (z0 + z1) : e, s)
exec Sub (ell, Left z0 : Left z1 : e, s)
                            = (ell+1, Left (z0 - z1) : e, s)
exec Mul (ell, Left z0 : Left z1 : e, s)
                            = (ell+1, Left (z0 * z1) : e, s)
exec PushTrue  (ell, e, s) = (ell+1, Right True  : e, s)
exec PushFalse (ell, e, s) = (ell+1, Right False : e, s)
exec Eq  (ell, Left z0 : Left z1 : e, s)
                            = (ell+1, Right (z0 == z1) : e, s)
exec Le  (ell, Left z0 : Left z1 : e, s)
                            = (ell+1, Right (z0 <= z1) : e, s)
exec Neg (ell, Right z : e, s)
                            = (ell+1, Right (not z) : e, s)
exec And (ell, Right z0 : Right z1 : e, s)
                            = (ell+1, Right (z0 && z1) : e, s)
exec Or  (ell, Right z0 : Right z1 : e, s)
                            = (ell+1, Right (z0 || z1) : e, s)
exec (Fetch x) (ell, e, s) = (ell+1, Left (lkp x s) : e, s)
exec (Store x) (ell, Left z : e, s) = (ell+1, e, upd x z s)
exec (Jump m) (ell, e, s) = (m, e, s)
exec (JumpFalse m) (ell, Right True  : e, s) = (ell+1, e, s)
exec (JumpFalse m) (ell, Right False : e, s) = (m, e, s)
  
 

run :: Code -> Config -> Config
run c cfg = case makeStep c cfg of
   Just cfg' -> run c cfg' 
   Nothing -> cfg