From iris.base_logic Require Import lib.iprop.
Module Type uses_spec.
    Parameter uses : Set.
    Parameter prop : ∀ {Σ}, uses -> iProp Σ -> iProp Σ.
    Parameter one : uses.
    Axiom one_spec : ∀ {Σ} {P : iProp Σ}, prop one P ⊣⊢ P.
    Parameter sep : uses -> uses -> uses.
    Axiom sep_spec : ∀ {Σ} {P : iProp Σ} u v, 
        prop (sep u v) P ⊢ prop u P ∗ prop v P.
    Parameter and : uses -> uses -> uses.
    Axiom and_spec : ∀ {Σ} {P : iProp Σ} u v, 
        prop (and u v) P ⊢ prop u P ∧ prop v P.
End uses_spec.
Module uses <: uses_spec.
    Inductive usesT := One | Many.
    Definition uses := usesT.
    Definition prop {Σ} (u : uses) (P : iProp Σ) : iProp Σ := match u with
        | One => P
        | Many => □P
    end.
    Definition one := One.
    Theorem one_spec {Σ} {P : iProp Σ} : prop one P ⊣⊢ P.
    Proof. done. Qed.    
    Definition sep (u : uses) (v : uses) : uses := Many.
    Theorem sep_spec {Σ} {P : iProp Σ} u v
        : prop (sep u v) P ⊢ prop u P ∗ prop v P.
    Proof.
        have: ∀ x, □P ⊢ prop x P by case; [apply: bi.intuitionistically_elim | ].
        move=> lemma.
        simpl.
        rewrite bi.intuitionistically_sep_dup.
        apply: bi.sep_mono; apply: lemma.
    Qed.
    Definition and (u : uses) (v : uses) : uses := match u, v with
        | One, One => One
        | _, _ => Many
    end.
    Theorem and_spec {Σ} {P : iProp Σ} u v
        : prop (and u v) P ⊢ prop u P ∧ prop v P.
    Proof.
        apply: bi.and_intro; case u; case v; try done; apply: bi.intuitionistically_elim.
    Qed.
End uses.

From stdpp Require Import binders strings stringmap.

From stdpp Require Import strings.

(* Set Implicit Arguments. *)
(* Nicer error messages. Instead of being asked to prove `False`,
    you are asked to prove `error (hopefully useful message)`.
*)
Module Type errorT.
    Parameter error : ∀ {T : Type}, T -> Prop.
    Axiom error_elim : ∀ {msgType} {msg : msgType} {T}, error msg -> T.
End errorT.
Module Export error : errorT.
    Definition error {T} (msg : T) := False.
    Theorem error_elim {msgType} (msg : msgType) T : error msg -> T. Proof. elim. Qed.
End error.

Require uses context expr.

Module context := context.context uses.uses.
Module Export expr := (expr.expr context).

(* TODO: Should I put this in a module, with a `:>`? *)
Inductive uses := One | Many.
Definition uses_prop {Σ} (u : uses) (P : iProp Σ) : iProp Σ := match u with
    | One => P
    | Many => □P
end.
Definition uses_sep (u : uses) (v : uses) : uses := Many.
Lemma uses_sep_spec {Σ} {P : iProp Σ} u v
    : uses_prop (uses_sep u v) P ⊢ uses_prop u P ∗ uses_prop v P.
Proof.
    have: ∀ x, □P ⊢ uses_prop x P by case; [apply: bi.intuitionistically_elim | ].
    move=> lemma.
    simpl.
    rewrite bi.intuitionistically_sep_dup.
    apply: bi.sep_mono; apply: lemma.
Qed.
Definition uses_and (u : uses) (v : uses) : uses := match u, v with
    | One, One => One
    | _, _ => Many
end.
Lemma uses_and_spec {Σ} {P : iProp Σ} u v
    : uses_prop (uses_and u v) P ⊢ uses_prop u P ∧ uses_prop v P.
Proof.
    apply: bi.and_intro; case u; case v; try done; apply: bi.intuitionistically_elim.
Qed.
Module Type contextT.
    (* Describes what must hold for an expression to typecheck.
        Includes:
        - Free variables, associated with their types and usage patterns.
        - Equality constraints, generated during typechecking.
    *)
    Parameter context : gFunctors -> Type.
    Parameter prop : ∀ {Σ}, context Σ -> iProp Σ.
    Parameter var : ∀ {Σ}, string -> iProp Σ -> context Σ.
    Axiom var_spec : ∀ {Σ} v (T : iProp Σ), prop (var v T) ⊢ T.
    Parameter constrain : ∀ {Σ}, Prop -> context Σ -> context Σ.
    Axiom constrain_spec : ∀ {Σ} P (ctx : context Σ),
        prop (constrain P ctx) ⊢ ∃ _ : P, prop ctx.
    (* Combine two expressions, in a separating way. *)
    Parameter sep : ∀ {Σ}, context Σ -> context Σ -> context Σ.
    Axiom sep_spec : ∀ {Σ} (c1 c2 : context Σ),
        prop (sep c1 c2) ⊢ prop c1 ∗ prop c2.
    
    (* Combine two expressions, in a non-separating way. *)
    Parameter and : ∀ {Σ}, context Σ -> context Σ -> context Σ.
    Axiom and_spec : ∀ {Σ} (c1 c2 : context Σ),
        prop (and c1 c2) ⊢ prop c1 ∧ prop c2.
    (* Remove a free variable from the context.
        Used to implement variable patterns.
    *)
    Parameter bind : ∀ {Σ}, string -> context Σ -> context Σ.
    (* Get the type and usage pattern of a variable. *)
    Parameter lookup : ∀ {Σ}, string -> context Σ -> option (iProp Σ * uses).
    
    Axiom bind_spec : ∀ {Σ} var (c : context Σ),
        match lookup var c with
            | None => emp
            | Some (P, uses) => (uses_prop uses P)
        end ∗ prop (bind var c) ⊢ 
        prop c.
    (* Interpret the expression in the toplevel scope.
        This will return a collection of constraints.
        Note that the constraint may be something like
        (error ("Variables not in scope.", ["x", "y"])).
    *)
    Parameter toplevel : ∀ {Σ}, context Σ -> Prop.
    Axiom toplevel_spec : ∀ {Σ} (ctx : context Σ),
        toplevel ctx -> ⊢ prop ctx.
End contextT.
Module context <: contextT.
    (* It turns out to be easier to carry the equality constraints 
        along with the variables, until they are bound.
        Afterward, I store the constraints separately.
    *)
    Record contextT {Σ} : Type := {
        Constraints : Prop;
        Variables : stringmap (iProp Σ * uses * Prop);
    }.
    Definition context Σ : Type := contextT (Σ := Σ).
    Definition prop {Σ} (ctx : context Σ) : iProp Σ := 
        ∃ _ : Constraints ctx,
        [∗ map] data ∈ Variables ctx, 
            let '(P, u, constraints) := data in 
            ∃ _ : (constraints : Prop), uses_prop u P.
    Definition var {Σ} (var : string) (type : iProp Σ) : context Σ := {|
        Constraints := True;
        Variables := {[var := (type, One, True)]};
    |}.
    Theorem var_spec {Σ} v (T : iProp Σ) : prop (var v T) ⊢ T.
    Proof.
        apply: bi.exist_elim; move=> [].
        rewrite big_sepM_singleton.
        apply: bi.exist_elim; move=> [].
        done.
    Qed.
    Definition constrain {Σ} (constraint : Prop) (ctx : context Σ) : context Σ := {|
        Constraints := constraint ∧ Constraints ctx;
        Variables := Variables ctx;
    |}.
    Theorem constrain_spec {Σ} P (ctx : context Σ)
        : prop (constrain P ctx) ⊢ ∃ _ : P, prop ctx.
    Proof.
        apply: bi.exist_elim. move=> [wP w].
        do 2 apply: bi.exist_intro'.
        done.
    Qed.
        
    Lemma big_sepM_hetero {type : bi} `{countable : Countable K} {V1 V2} 
        (Φ : K -> V1 -> type) (Ψ : K -> V2 -> type)
        (m1 : gmap K V1) (m2 : gmap K V2)
        : (∀ k, 
            (if m1 !! k is Some v then Φ k v else emp) ⊢ 
            (if m2 !! k is Some v then Ψ k v else emp))
        -> ([∗ map] k ↦ v ∈ m1, Φ k v) ⊢ ([∗ map] k ↦ v ∈ m2, Ψ k v).
    Proof.
        move: m2.
        induction m1 as [|k v1 m1' fact1 IH1] using map_ind.
        - move=> m2 H. rewrite big_sepM_empty.
            induction m2 as [|k v2 m2' fact2 IH2] using map_ind.
            + by rewrite big_sepM_empty.
            + rewrite big_sepM_insert; last exact fact2.
                rewrite -(bi.emp_sep emp%I).
                apply: bi.sep_mono.
                * move: (H k) => H'.
                    by rewrite lookup_empty lookup_insert in H'.
                * apply: IH2 => k'.
                    rewrite lookup_empty.
                    case: (decide (k = k')).
                    -- move=><-. by rewrite fact2.
                    -- move=> neq.
                        move: (H k').
                        by rewrite lookup_empty lookup_insert_ne.
        - move=> m2 H.
            rewrite big_sepM_insert; last exact fact1.
            apply: transitivity.
            { 
                apply: bi.sep_mono.
                - move: (H k) => H'.
                    rewrite lookup_insert in H'.
                    exact: H'.
                - apply: (IH1 (delete k m2)) => k'.
                    case: (decide (k = k')).
                    + move=><-. by rewrite fact1 lookup_delete.
                    + move=> neq.
                        rewrite lookup_delete_ne; last exact neq.
                        move: (H k').
                        rewrite lookup_insert_ne; last exact neq.
                        done.
            }
    
            have: m2 !! k = m2 !! k by []; case: {-1}(m2 !! k) => [v2|] fact2.
            + by rewrite (big_sepM_delete _ m2).
            + by rewrite bi.emp_sep delete_notin.
    Qed.
    Definition sep {Σ} (c1 : context Σ) (c2 : context Σ) : context Σ := {| 
        Constraints := Constraints c1 ∧ Constraints c2;
        Variables := union_with
            (λ '(P, u, c1) '(Q, v, c2), Some (P, uses_sep u v, c1 ∧ c2 ∧ P = Q)) 
            (Variables c1) (Variables c2);
    |}.
    Theorem sep_spec {Σ} (c1 c2 : context Σ)
        : prop (sep c1 c2) ⊢ prop c1 ∗ prop c2.
    Proof.
        rewrite /prop.
        apply: bi.exist_elim; move => [w1 w2].
        rewrite bi.sep_exist_l; apply: bi.exist_intro'.
        rewrite bi.sep_exist_r; apply: bi.exist_intro'.
        clear w1 w2.
        
        apply: transitivity; last (
            apply: bi.sep_mono;
            [ apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c1 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses_prop u P 
                    else emp)%I _ 
                (Variables (sep c1 c2))) => k
            | apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c2 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses_prop u P 
                    else emp)%I _ 
                (Variables (sep c1 c2))) => k
            ]
        ).
        - rewrite -big_sepM_sep.
            apply: big_sepM_mono => k [[P u] constraints] fact.
            apply: bi.exist_elim.
            move: fact; rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            simpl; last done;
            move=>w; injection w; clear w;
            move=> <-<-<- w.
            + move: w => [w1 [w2 <-]].
                rewrite bi.sep_exist_l; apply: bi.exist_intro'.
                rewrite bi.sep_exist_r; apply: bi.exist_intro'.
                apply: uses_sep_spec.
            + by rewrite bi.sep_emp; apply: bi.exist_intro'.
            + by rewrite bi.emp_sep; apply: bi.exist_intro'.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
    Qed.
    Definition and {Σ} (c1 : context Σ) (c2 : context Σ) : context Σ := {| 
        Constraints := Constraints c1 ∧ Constraints c2;
        Variables := union_with
            (λ '(P, u, c1) '(Q, v, c2), Some (P, uses_and u v, c1 ∧ c2 ∧ P = Q)) 
            (Variables c1) (Variables c2);
    |}.
    Theorem and_spec {Σ} (c1 c2 : context Σ)
        : prop (and c1 c2) ⊢ prop c1 ∧ prop c2.
    Proof.
        rewrite /prop.
        apply: bi.exist_elim; move => [w1 w2].
        rewrite bi.and_exist_l; apply: bi.exist_intro'.
        rewrite bi.and_exist_r; apply: bi.exist_intro'.
        clear w1 w2.
        apply: transitivity; last (
            apply: bi.and_mono;
            [ apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c1 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses_prop u P 
                    else emp)%I _ 
                (Variables (and c1 c2))) => k
            | apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c2 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses_prop u P 
                    else emp)%I _ 
                (Variables (and c1 c2))) => k
            ]
        ).
        - rewrite -big_sepM_and. 
            apply: big_sepM_mono => k [[P u] constraints] fact.
            apply: bi.exist_elim.
            move: fact; rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            simpl; last done;
            move=>w; injection w; clear w;
            move=> <-<-<- w.
            + move: w => [w1 [w2 <-]].
                rewrite bi.and_exist_l; apply: bi.exist_intro'.
                rewrite bi.and_exist_r; apply: bi.exist_intro'.
                apply: uses_and_spec.
            + by rewrite bi.and_True; apply: bi.exist_intro'.
            + by rewrite bi.True_and; apply: bi.exist_intro'.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
    Qed.
    Definition bind {Σ} (var : string) (ctx : context Σ) : context Σ := {|
        Variables := delete var (Variables ctx);
        Constraints := Constraints ctx ∧ 
            if Variables ctx !! var is Some (_, _, constraints)
            then constraints 
            else True;
    |}.
    Definition lookup {Σ} (var : string) (ctx : context Σ)
        : option (iProp Σ * uses)
    := option_map fst (Variables ctx !! var).
    Theorem bind_spec {Σ} var (c : context Σ)
        : match lookup var c with
            | None => emp
            | Some (P, uses) => (uses_prop uses P)
        end ∗ prop (bind var c) ⊢ 
        prop c.
    Proof.
        rewrite /prop.
        rewrite bi.sep_exist_l; apply: bi.exist_elim.
        move=> [w1 w2].
        apply: bi.exist_intro'.
        rewrite /lookup.
        move: w2.
        have: Variables c !! var = Variables c !! var by [];
            case: {-1}(Variables c !! var) => [[[P u] constraints]|] fact w2.
        - simpl.
            rewrite (big_sepM_delete _ (Variables c)); last exact fact.
            rewrite bi.sep_exist_r. apply: bi.exist_intro'.
            done.
        - simpl. 
            by rewrite bi.emp_sep delete_notin.
    Qed.
    Definition toplevel {Σ} (ctx : context Σ) : Prop :=
        if decide (Variables ctx = ∅) 
        then Constraints ctx
        else error (
            "Variables not in scope.", 
            map fst (map_to_list (Variables ctx))
        ).
    Theorem toplevel_spec {Σ} (ctx : context Σ) :
        toplevel ctx -> ⊢ prop ctx.
    Proof.
        rewrite /toplevel /prop.
        case: (decide (Variables ctx = ∅)) => [-> | _] w.
        - apply: bi.exist_intro'.
            by rewrite big_sepM_empty.
        - by apply: error_elim.
    Qed.
End context.
Inductive expr {Σ : gFunctors} : iProp Σ -> Type :=
    | Var {T} : string -> expr T
    | Lam {T1 T2} : string -> expr T2 -> expr (T1 -∗ T2)
    | Pair {T1 T2} : expr T1 -> expr T2 -> expr (T1 ∗ T2)
.
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 = 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.
(* Example Proof. *)
Goal ∀ {Σ} (P Q : iProp Σ), ⊢ P -∗ P.
    move=> Σ P Q.
    apply: (check_program_spec (Lam "x" (Var "x"))).
    by vm_compute.
Qed.
(* Example Proof. *)
Goal ∀ {Σ} (P Q : iProp Σ), ⊢ P -∗ Q -∗ Q ∗ P.
    move=> Σ P Q.
    apply: (check_program_spec (Lam "x" (Lam "y" (Pair (Var "y") (Var "x"))))).
    by vm_compute.
Qed.
(* Example Proof. *)
Goal ∀ {Σ} (P Q R : iProp Σ), ⊢ P -∗ Q -∗ R -∗ (R ∗ Q) ∗ P.
    move=> Σ P Q R.
    apply: (check_program_spec (Lam "x" (Lam "y" (Lam "z" (Pair (Pair (Var "z") (Var "y")) (Var "x")))))).
    by vm_compute.
Qed.

Export uses.uses.

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.

(* Nicer error messages. Instead of being asked to prove `False`,
    you are asked to prove `error (hopefully useful message)`.
*)
Module Type error_spec.
    Parameter error : forall {T : Type}, T -> Prop.
    Axiom error_elim : forall {msgType} {msg : msgType} {T}, error msg -> T.
End error_spec.
Module Export error : error_spec.
    Definition error {T} (msg : T) := False.
    Definition error_elim {msgType} (msg : msgType) T : error msg -> T := False_rect T.
End error.

From stdpp Require Import binders strings stringmap.
From iris.base_logic Require Import lib.iprop.
Require Import error uses.
Module Type context_spec (uses : uses_spec).
    (* Describes what must hold for an expression to typecheck.
        Includes:
        - Free variables, associated with their types and usage patterns.
        - Equality constraints, generated during typechecking.
    *)
    Parameter context : gFunctors -> Type.
    Parameter prop : ∀ {Σ}, context Σ -> iProp Σ.
    Parameter var : ∀ {Σ}, string -> iProp Σ -> context Σ.
    Axiom var_spec : ∀ {Σ} v (T : iProp Σ), prop (var v T) ⊢ T.
    Parameter constrain : ∀ {Σ}, Prop -> context Σ -> context Σ.
    Axiom constrain_spec : ∀ {Σ} P (ctx : context Σ),
        prop (constrain P ctx) ⊢ ∃ _ : P, prop ctx.
    (* Combine two expressions, in a separating way. *)
    Parameter sep : ∀ {Σ}, context Σ -> context Σ -> context Σ.
    Axiom sep_spec : ∀ {Σ} (c1 c2 : context Σ),
        prop (sep c1 c2) ⊢ prop c1 ∗ prop c2.
    
    (* Combine two expressions, in a non-separating way. *)
    Parameter and : ∀ {Σ}, context Σ -> context Σ -> context Σ.
    Axiom and_spec : ∀ {Σ} (c1 c2 : context Σ),
        prop (and c1 c2) ⊢ prop c1 ∧ prop c2.
    (* Remove a free variable from the context.
        Used to implement variable patterns.
    *)
    Parameter bind : ∀ {Σ}, string -> context Σ -> context Σ.
    (* Get the type and usage pattern of a variable. *)
    Parameter lookup : ∀ {Σ}, string -> context Σ -> option (iProp Σ * uses.uses).
    
    Axiom bind_spec : ∀ {Σ} var (c : context Σ),
        match lookup var c with
            | None => emp
            | Some (P, uses) => (uses.prop uses P)
        end ∗ prop (bind var c) ⊢ 
        prop c.
    (* Interpret the expression in the toplevel scope.
        This will return a collection of constraints.
        Note that the constraint may be something like
        (error ("Variables not in scope.", ["x", "y"])).
    *)
    Parameter toplevel : ∀ {Σ}, context Σ -> Prop.
    Axiom toplevel_spec : ∀ {Σ} (ctx : context Σ),
        toplevel ctx -> ⊢ prop ctx.
End context_spec.
Module context (uses : uses_spec) <: context_spec(uses).
    (* It turns out to be easier to carry the equality constraints 
        along with the variables, until they are bound.
        Afterward, I store the constraints separately.
    *)
    Record contextT {Σ} : Type := {
        Constraints : Prop;
        Variables : stringmap (iProp Σ * uses.uses * Prop);
    }.
    Definition context Σ : Type := contextT (Σ := Σ).
    Definition prop {Σ} (ctx : context Σ) : iProp Σ := 
        ∃ _ : Constraints ctx,
        [∗ map] data ∈ Variables ctx, 
            let '(P, u, constraints) := data in 
            ∃ _ : (constraints : Prop), uses.prop u P.
    Definition var {Σ} (var : string) (type : iProp Σ) : context Σ := {|
        Constraints := True;
        Variables := {[var := (type, uses.one, True)]};
    |}.
    Theorem var_spec {Σ} v (T : iProp Σ) : prop (var v T) ⊢ T.
    Proof.
        apply: bi.exist_elim; move=> [].
        rewrite big_sepM_singleton.
        apply: bi.exist_elim; move=> [].
        by rewrite uses.one_spec.
    Qed.
    Definition constrain {Σ} (constraint : Prop) (ctx : context Σ) : context Σ := {|
        Constraints := constraint ∧ Constraints ctx;
        Variables := Variables ctx;
    |}.
    Theorem constrain_spec {Σ} P (ctx : context Σ)
        : prop (constrain P ctx) ⊢ ∃ _ : P, prop ctx.
    Proof.
        apply: bi.exist_elim. move=> [wP w].
        do 2 apply: bi.exist_intro'.
        done.
    Qed.
        
    Lemma big_sepM_hetero {type : bi} `{countable : Countable K} {V1 V2} 
        (Φ : K -> V1 -> type) (Ψ : K -> V2 -> type)
        (m1 : gmap K V1) (m2 : gmap K V2)
        : (∀ k, 
            (if m1 !! k is Some v then Φ k v else emp) ⊢ 
            (if m2 !! k is Some v then Ψ k v else emp))
        -> ([∗ map] k ↦ v ∈ m1, Φ k v) ⊢ ([∗ map] k ↦ v ∈ m2, Ψ k v).
    Proof.
        move: m2.
        induction m1 as [|k v1 m1' fact1 IH1] using map_ind.
        - move=> m2 H. rewrite big_sepM_empty.
            induction m2 as [|k v2 m2' fact2 IH2] using map_ind.
            + by rewrite big_sepM_empty.
            + rewrite big_sepM_insert; last exact fact2.
                rewrite -(bi.emp_sep emp%I).
                apply: bi.sep_mono.
                * move: (H k) => H'.
                    by rewrite lookup_empty lookup_insert in H'.
                * apply: IH2 => k'.
                    rewrite lookup_empty.
                    case: (decide (k = k')).
                    -- move=><-. by rewrite fact2.
                    -- move=> neq.
                        move: (H k').
                        by rewrite lookup_empty lookup_insert_ne.
        - move=> m2 H.
            rewrite big_sepM_insert; last exact fact1.
            apply: transitivity.
            { 
                apply: bi.sep_mono.
                - move: (H k) => H'.
                    rewrite lookup_insert in H'.
                    exact: H'.
                - apply: (IH1 (delete k m2)) => k'.
                    case: (decide (k = k')).
                    + move=><-. by rewrite fact1 lookup_delete.
                    + move=> neq.
                        rewrite lookup_delete_ne; last exact neq.
                        move: (H k').
                        rewrite lookup_insert_ne; last exact neq.
                        done.
            }
    
            have: m2 !! k = m2 !! k by []; case: {-1}(m2 !! k) => [v2|] fact2.
            + by rewrite (big_sepM_delete _ m2).
            + by rewrite bi.emp_sep delete_notin.
    Qed.
    Definition sep {Σ} (c1 : context Σ) (c2 : context Σ) : context Σ := {| 
        Constraints := Constraints c1 ∧ Constraints c2;
        Variables := union_with
            (λ '(P, u, c1) '(Q, v, c2), Some (P, uses.sep u v, c1 ∧ c2 ∧ P = Q)) 
            (Variables c1) (Variables c2);
    |}.
    Theorem sep_spec {Σ} (c1 c2 : context Σ)
        : prop (sep c1 c2) ⊢ prop c1 ∗ prop c2.
    Proof.
        rewrite /prop.
        apply: bi.exist_elim; move => [w1 w2].
        rewrite bi.sep_exist_l; apply: bi.exist_intro'.
        rewrite bi.sep_exist_r; apply: bi.exist_intro'.
        clear w1 w2.
        
        apply: transitivity; last (
            apply: bi.sep_mono;
            [ apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c1 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses.prop u P 
                    else emp)%I _ 
                (Variables (sep c1 c2))) => k
            | apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c2 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses.prop u P 
                    else emp)%I _ 
                (Variables (sep c1 c2))) => k
            ]
        ).
        - rewrite -big_sepM_sep.
            apply: big_sepM_mono => k [[P u] constraints] fact.
            apply: bi.exist_elim.
            move: fact; rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            simpl; last done;
            move=>w; injection w; clear w;
            move=> <-<-<- w.
            + move: w => [w1 [w2 <-]].
                rewrite bi.sep_exist_l; apply: bi.exist_intro'.
                rewrite bi.sep_exist_r; apply: bi.exist_intro'.
                apply: uses.sep_spec.
            + by rewrite bi.sep_emp; apply: bi.exist_intro'.
            + by rewrite bi.emp_sep; apply: bi.exist_intro'.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
    Qed.
    Definition and {Σ} (c1 : context Σ) (c2 : context Σ) : context Σ := {| 
        Constraints := Constraints c1 ∧ Constraints c2;
        Variables := union_with
            (λ '(P, u, c1) '(Q, v, c2), Some (P, uses.and u v, c1 ∧ c2 ∧ P = Q)) 
            (Variables c1) (Variables c2);
    |}.
    Theorem and_spec {Σ} (c1 c2 : context Σ)
        : prop (and c1 c2) ⊢ prop c1 ∧ prop c2.
    Proof.
        rewrite /prop.
        apply: bi.exist_elim; move => [w1 w2].
        rewrite bi.and_exist_l; apply: bi.exist_intro'.
        rewrite bi.and_exist_r; apply: bi.exist_intro'.
        clear w1 w2.
        apply: transitivity; last (
            apply: bi.and_mono;
            [ apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c1 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses.prop u P 
                    else emp)%I _ 
                (Variables (and c1 c2))) => k
            | apply: (big_sepM_hetero
                (λ k _, 
                    if Variables c2 !! k is Some (P, u, constraints) 
                    then ∃ _ : (constraints : Prop), uses.prop u P 
                    else emp)%I _ 
                (Variables (and c1 c2))) => k
            ]
        ).
        - rewrite -big_sepM_and. 
            apply: big_sepM_mono => k [[P u] constraints] fact.
            apply: bi.exist_elim.
            move: fact; rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            simpl; last done;
            move=>w; injection w; clear w;
            move=> <-<-<- w.
            + move: w => [w1 [w2 <-]].
                rewrite bi.and_exist_l; apply: bi.exist_intro'.
                rewrite bi.and_exist_r; apply: bi.exist_intro'.
                apply: uses.and_spec.
            + by rewrite bi.and_True; apply: bi.exist_intro'.
            + by rewrite bi.True_and; apply: bi.exist_intro'.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
        - rewrite lookup_union_with.
            case: (Variables c1 !! k) => [[[P1 u1] constraints1]|];
            case: (Variables c2 !! k) => [[[P2 u2] constraints2]|];
            done.
    Qed.
    Definition bind {Σ} (var : string) (ctx : context Σ) : context Σ := {|
        Variables := delete var (Variables ctx);
        Constraints := Constraints ctx ∧ 
            if Variables ctx !! var is Some (_, _, constraints)
            then constraints 
            else True;
    |}.
    Definition lookup {Σ} (var : string) (ctx : context Σ)
        : option (iProp Σ * uses.uses)
    := option_map fst (Variables ctx !! var).
    Theorem bind_spec {Σ} var (c : context Σ)
        : match lookup var c with
            | None => emp
            | Some (P, uses) => (uses.prop uses P)
        end ∗ prop (bind var c) ⊢ 
        prop c.
    Proof.
        rewrite /prop.
        rewrite bi.sep_exist_l; apply: bi.exist_elim.
        move=> [w1 w2].
        apply: bi.exist_intro'.
        rewrite /lookup.
        move: w2.
        have: Variables c !! var = Variables c !! var by [];
            case: {-1}(Variables c !! var) => [[[P u] constraints]|] fact w2.
        - simpl.
            rewrite (big_sepM_delete _ (Variables c)); last exact fact.
            rewrite bi.sep_exist_r. apply: bi.exist_intro'.
            done.
        - simpl. 
            by rewrite bi.emp_sep delete_notin.
    Qed.
    Definition toplevel {Σ} (ctx : context Σ) : Prop :=
        if decide (Variables ctx = ∅) 
        then Constraints ctx
        else error (
            "Variables not in scope.", 
            map fst (map_to_list (Variables ctx))
        ).
    Theorem toplevel_spec {Σ} (ctx : context Σ) :
        toplevel ctx -> ⊢ prop ctx.
    Proof.
        rewrite /toplevel /prop.
        case: (decide (Variables ctx = ∅)) => [-> | _] w.
        - apply: bi.exist_intro'.
            by rewrite big_sepM_empty.
        - by apply: error_elim.
    Qed.
End context.

From stdpp Require Import binders strings stringmap.
From iris.base_logic Require Import lib.iprop.
Require Import lang.
(* Example Proof. *)
Goal ∀ {Σ} (P Q : iProp Σ), ⊢ P -∗ P.
    move=> Σ P Q.
    apply: (check_program_spec (Lam "x" (Var "x"))).
    by vm_compute.
Qed.

theories/error.v
theories/uses.v
theories/context.v
theories/expr.v

examples/basic.v