From stdpp Require Import strings.
From iris.base_logic Require Import lib.iprop.

Require uses context.
Import uses.

Module Type expr_spec.
    Parameter expr : ∀ {Σ}, iProp Σ -> Type.

    Parameter check_program : ∀ {Σ} {type : iProp Σ}, expr type -> Prop.

    Axiom check_program_spec : ∀ {Σ} {type : iProp Σ} (e : expr type),
        check_program e -> ⊢ type.
End expr_spec.

Module expr (context : context.context_spec uses).
    Inductive exprT {Σ : gFunctors} : iProp Σ -> Type :=
        | Var {T} : string -> exprT T
        | Lam {T1 T2} : string -> exprT T2 -> exprT (T1 -∗ T2)
        | Pair {T1 T2} : exprT T1 -> exprT T2 -> exprT (T1 ∗ T2)
    .
    Definition expr {Σ} := exprT (Σ := Σ).

    Fixpoint check_expr {Σ} {type : iProp Σ} (e : expr type) : context.context Σ := match e with
        | Var var => context.var var type
        | @Lam _ T1 T2 var e =>
            let ctx := check_expr e in
            context.constrain
                (match context.lookup var ctx with
                    | Some (T, uses) => T = T1 ∧ uses = uses.One
                    | None => True
                end)
                (context.bind var (check_expr e))
        | Pair e1 e2 => context.sep (check_expr e1) (check_expr e2)
    end.
    
    Lemma check_expr_spec {Σ} (type : iProp Σ) (e : expr type) : 
        context.prop (check_expr e) ⊢ type.
    Proof.
        elim e; simpl.
        - move=> T var. apply: context.var_spec.
        - move=> T1 T2 var body H.
            apply: transitivity; first apply: context.constrain_spec.
            apply: bi.exist_elim => w.

            apply: bi.wand_intro_r.
            apply: transitivity; last apply: H.
            apply: transitivity; last apply: (context.bind_spec var).

            rewrite {1}bi.sep_comm.
            apply: bi.sep_mono_l.

            move: w; case: (context.lookup var (check_expr body)) => [[P u]|].
            + by move=>[->->].
            + move=>[]. apply: bi.True_intro.
        - move=> T1 T2 e1 H1 e2 H2.
            apply: transitivity; first apply: context.sep_spec.
            by apply: bi.sep_mono.
    Qed.

    Definition check_program {Σ} {type : iProp Σ} (e : expr type) : Prop :=
        context.toplevel (check_expr e).

    Theorem check_program_spec {Σ} {type : iProp Σ} (e : expr type) : 
        check_program e -> ⊢ type.
    Proof.
        move=> w.
        apply: transitivity.
        - apply: context.toplevel_spec w.
        - apply: check_expr_spec.
    Qed.
End expr.