; type checking for simply typed lambda calculus

(datatype Type 
  (TArr Type Type) ; t1 -> t2
)
(function TUnitConst () Type)
(let TUnit (TUnitConst))

(datatype Expr 
  (Lam String Type Expr) ; lam x : t . e
  (App Expr Expr) 
  (Var String) 
)
(function MyUnitConst () Expr)
(let MyUnit (MyUnitConst))

(datatype Ctx 
  (Cons String Type Ctx)
)
(function NilConst () Ctx)
(let Nil (NilConst))

; ctx |- expr : type
(function typeof (Ctx Expr) Type)

; ctx |- () : unit
(rewrite (typeof ctx MyUnit) TUnit)

; ctx; x: t |- x : t
(rewrite (typeof (Cons x t ctx) (Var x)) t)

; ctx |- f :- t1 -> t2
; ctx |- e : t1
; -----------------
; ctx |- f e : t2

(rule (
  (= (typeof ctx (App f e)) t2)
)(
  (typeof ctx f)
  (typeof ctx e)
))

(rule (
  (= (typeof ctx (App f e)) t1)
  (= (typeof ctx f) (TArr (typeof ctx e) t2))
)(
  (union t1 t2)
))

; ctx |- x : t
; ------------------ y != x 
; ctx; y: t |- x : t

(rewrite (typeof (Cons y ty ctx) (Var x))
         (typeof ctx (Var x))
    :when ((!= x y)))

; ctx; x: t1 |- e : t2
; ------------------------------
; ctx |- lam x: t1. e : t1 -> t2

; rhs of rewrite creates demand
(rewrite (typeof ctx (Lam x t1 e))
         (TArr t1 (typeof (Cons x t1 ctx) e)))

; TEST
; ----

; lam x : unit, f : unit -> unit . f x
(let e 
  (Lam "x" TUnit 
       (Lam "f" (TArr TUnit TUnit)
            (App (Var "f") (Var "x")))))

; lam x : unit . x
(let id (Lam "x" TUnit (Var "x")))
(let t-id (typeof Nil id))

; (e () id) = ()
(let app-unit-id (App (App e MyUnit) id))
(let t-app (typeof Nil app-unit-id))

(let free (Lam "x" TUnit (Var "y")))
(let t-free-ill (typeof Nil free))
(let t-free-1 (typeof (Cons "y" TUnit Nil) free))
(let t-free-2 (typeof (Cons "y" (TArr (TArr TUnit TUnit) TUnit) Nil) free))

(run 15)

(query-extract t-id)
(check (= t-id (TArr TUnit TUnit)))

(query-extract t-app)
(check (= t-app TUnit))

(query-extract t-free-1)
(check (= t-free-1 (TArr TUnit TUnit)))
(query-extract t-free-2)
(check (= t-free-2 (TArr TUnit (TArr (TArr TUnit TUnit) TUnit))))
; this will err
; (query-extract t-free-ill)