data Ty = IntTy | BoolTy
data Reg = A | B | R

StateTy : Type
StateTy = (Ty, Ty, Ty)

getRegTy : Reg -> StateTy -> Ty
getRegTy A (a, _, _) = a
getRegTy B (_, b, _) = b
getRegTy R (_, _, r) = r

setRegTy : Reg -> Ty -> StateTy -> StateTy
setRegTy A a (_, b, r) = (a, b, r)
setRegTy B b (a, _, r) = (a, b, r)
setRegTy R r (a, b, _) = (a, b, r)

data Expr : StateTy -> Ty -> Type where
  Lit : Int -> Expr s IntTy
  Load : (r : Reg) -> Expr s (getRegTy r s)
  Add : Expr s IntTy -> Expr s IntTy -> Expr s IntTy
  Leq : Expr s IntTy -> Expr s IntTy -> Expr s BoolTy
  Not : Expr s BoolTy -> Expr s BoolTy

mutual
  data Stmt : StateTy -> StateTy -> Type where
    Store : (r : Reg) -> Expr s t -> Stmt s (setRegTy r t s)
    If : Expr s BoolTy -> Prog s n -> Prog s n -> Stmt s n
    Loop : Prog s s -> Stmt s s
    Break : Stmt s s
  
  data Prog : StateTy -> StateTy -> Type where
    Nil : Prog s s
    (::) : Stmt s n -> Prog n m -> Prog s m

initialState : StateTy
initialState = (IntTy, IntTy, IntTy)

testProg : Prog Main.initialState Main.initialState
testProg =
  [ Store A (Lit 1 `Leq` Lit 2)
  , If (Load A)
    [ Store A (Lit 1) ]
    [ Store A (Lit 2) ]
  , Store B (Lit 2)
  , Store R (Add (Load A) (Load B))
  ]