content: {-@ LIQUID "--exact-data-cons" @-}
  {-@ LIQUID "--ple" @-}
  
  module Rest where
  
  import Language.Haskell.Liquid.ProofCombinators
  import Prelude hiding (length, max)
  
  -- | The type of a Stitch expression
  data Ty = TInt
          | TBool
          | TFun { funArgTy :: Ty, funResTy :: Ty }
    deriving (Show, Eq)
  
  {-@
  data Ty = TInt
          | TBool
          | TFun { funArgTy :: Ty, funResTy :: Ty }
  @-}
  
  {-@ measure isFunTy @-}
  isFunTy :: Ty -> Bool
  isFunTy (TFun _ _) = True
  isFunTy _ = False
  
  
  -- | An @ArithOp ty@ is an operator on numbers that produces a result
  -- of type @ty@
  data ArithOp
    = Plus
    | Minus
    | Times
    | Divide
    | Mod
    | Less
    | LessE
    | Greater
    | GreaterE
    | Equals
    deriving (Eq, Show)
  
  -- | Extract the result type of an Op
  {-@ measure arithType @-}
  arithType :: ArithOp -> Ty
  arithType Plus     = TInt
  arithType Minus    = TInt
  arithType Times    = TInt
  arithType Divide   = TInt
  arithType Mod      = TInt
  arithType Less     = TBool
  arithType LessE    = TBool
  arithType Greater  = TBool
  arithType GreaterE = TBool
  arithType Equals   = TBool
  
  {-@ invariant { op:ArithOp | arithType op = TBool || arithType op = TInt } @-}
  
  
  type Nat = Int
  
  {-@ inline max @-}
  max :: Ord a => a -> a -> a
  max a b = if a > b then a else b
  
  {-@ inline minus @-}
  minus :: Nat -> Nat -> Nat
  minus a b = max 0 (a - b)
  
  {-@
  reflect elemAt
  elemAt :: n : Nat -> { xs:[a] | length xs > n } -> a
  @-}
  elemAt :: Nat -> [a] -> a
  elemAt 0 (x:_) = x
  elemAt i (_:xs) = elemAt (i-1) xs
  
  {-@
  measure length
  length :: xs : [a] -> Nat
  @-}
  length :: [a] -> Nat
  length [] = 0
  length (_:xs) = 1 + length xs
  
  {-@
  reflect append
  append ::
    xs : [a] ->
    ys : [a] ->
    { zs:[a] | length zs == length xs + length ys }
   @-}
  append :: [a] -> [a] -> [a]
  append [] ys = ys
  append (x:xs) ys = x : append xs ys
  
  {-@
  appendLengh
    :: xs : [a]
    -> ys : [a]
    -> { length (append xs ys) == length xs + length ys}
  @-}
  appendLengh :: [a] -> [a] -> ()
  appendLengh xs ys = () ? append xs ys
  
  {-@
  elemAtThroughAppend
    :: i : Nat
    -> xs : { xs:[a] | i < length xs }
    -> ys : [a]
    -> { elemAt i (append xs ys) = elemAt i xs }
  @-}
  elemAtThroughAppend :: Nat -> [a] -> [a] -> ()
  elemAtThroughAppend 0 (_:_) ys = ()
  elemAtThroughAppend i (_:xss) ys = elemAtThroughAppend (i - 1) xss ys
  
  {-@
  predicate WellTyped E CTX = checkBindings CTX E && numFreeVarsExp E <= length CTX
  type WellTypedExp CTX = { e : Exp | WellTyped e CTX }
  type FunExp = { e : Exp | isFunTy (exprType e) }
  type ExpT T = { e : Exp | T = exprType e }
  data Exp
    = Var Ty Nat
    | Lam Ty Exp
    | App (e1 :: FunExp) (ExpT (funArgTy (exprType e1)))
    | Let Exp Exp
    | Arith (ExpT TInt) ArithOp (ExpT TInt)
    | Cond (ExpT TBool) (a :: Exp) (ExpT (exprType a))
    | Fix ({ e:FunExp | funArgTy (exprType e) = funResTy (exprType e) })
    | IntE Int
    | BoolE Bool
  @-}
  
  -- | Checked expression
  data Exp
    = Var Ty Nat
    | Lam Ty Exp
    | App Exp Exp
    | Let Exp Exp
    | Arith Exp ArithOp Exp
    | Cond Exp Exp Exp
    | Fix Exp
    | IntE Int
    | BoolE Bool
    deriving Show
  
  {-@ measure exprType @-}
  exprType :: Exp -> Ty
  exprType (Var ty _) = ty
  exprType (Lam ty e) = TFun ty (exprType e)
  exprType (App e1 _) = funResTy (exprType e1)
  exprType (Let _ e2) = exprType e2
  exprType (Arith _ op _) = arithType op
  exprType (Cond _ e2 _) = exprType e2
  exprType (Fix e) = funResTy (exprType e)
  exprType (IntE _) = TInt
  exprType (BoolE _) = TBool
  
  -- | Check that all occurrences of a variable have the given type
  {-@
  reflect checkBindings
  checkBindings
    :: ctx : [Ty]
    -> { e : Exp | numFreeVarsExp e <= length ctx }
    -> Bool
  @-}
  checkBindings :: [Ty] -> Exp -> Bool
  checkBindings ctx (Var vty i) = elemAt i ctx == vty
  checkBindings ctx (Lam t e) = checkBindings (t:ctx) e
  checkBindings ctx (App e1 e2) = checkBindings ctx e1 && checkBindings ctx e2
  checkBindings ctx (Let e1 e2) = checkBindings ctx e1 && checkBindings (exprType e1 : ctx) e2
  checkBindings ctx (Arith e1 _ e2) = checkBindings ctx e1 && checkBindings ctx e2
  checkBindings ctx (Cond e1 e2 e3) = checkBindings ctx e1 && checkBindings ctx e2 && checkBindings ctx e3
  checkBindings ctx (Fix e) = checkBindings ctx e
  checkBindings _ (IntE _) = True
  checkBindings _ (BoolE _) = True
  
  {-@
  rewriteWith aClosedExpIsValidInAnyContext [appendLengh]
  aClosedExpIsValidInAnyContext
    :: ctx0 : [Ty]
    -> ctx1 : [Ty]
    -> e : Exp
    -> { WellTyped e ctx0 <=>
         WellTyped e (append ctx0 ctx1) && numFreeVarsExp e <= length ctx0
       }
  @-}
  aClosedExpIsValidInAnyContext :: [Ty] -> [Ty] -> Exp -> ()
  aClosedExpIsValidInAnyContext ctx0 ctx1 e = case e of
    Var _ i ->
      if i < length ctx0 then elemAtThroughAppend i ctx0 ctx1
      else ()
    Lam ty body ->
      aClosedExpIsValidInAnyContext (ty:ctx0) ctx1 body
    App e1 e2 ->
      aClosedExpIsValidInAnyContext ctx0 ctx1 e1 ? aClosedExpIsValidInAnyContext ctx0 ctx1 e2
    Let e1 e2 ->
      aClosedExpIsValidInAnyContext ctx0 ctx1 e1 ? aClosedExpIsValidInAnyContext (exprType e1 : ctx0) ctx1 e2
    Arith e1 _ e2 ->
      aClosedExpIsValidInAnyContext ctx0 ctx1 e1 ? aClosedExpIsValidInAnyContext ctx0 ctx1 e2
    Cond e1 e2 e3 ->
      aClosedExpIsValidInAnyContext ctx0 ctx1 e1
        ? aClosedExpIsValidInAnyContext ctx0 ctx1 e2
        ? aClosedExpIsValidInAnyContext ctx0 ctx1 e3
    Fix body ->
      aClosedExpIsValidInAnyContext ctx0 ctx1 body
    IntE _ -> ()
    BoolE _ -> ()
  
  {-@
  measure numFreeVarsExp
  numFreeVarsExp :: Exp -> Nat
  @-}
  numFreeVarsExp :: Exp -> Nat
  numFreeVarsExp (Var _ v) = v + 1
  numFreeVarsExp (Lam _ body) = minus (numFreeVarsExp body) 1
  numFreeVarsExp (App e1 e2) = max (numFreeVarsExp e1) (numFreeVarsExp e2)
  numFreeVarsExp (Let e1 e2) =
      max (numFreeVarsExp e1) (minus (numFreeVarsExp e2) 1)
  numFreeVarsExp (Arith e1 _ e2) = max (numFreeVarsExp e1) (numFreeVarsExp e2)
  numFreeVarsExp (Cond e1 e2 e3) =
      max (max (numFreeVarsExp e1) (numFreeVarsExp e2)) (numFreeVarsExp e3)
  numFreeVarsExp (Fix body) = numFreeVarsExp body
  numFreeVarsExp (IntE _) = 0
  numFreeVarsExp (BoolE _) = 0
  
  
  