Module context := context.context uses.uses.

Module context := context.context_simpl_impl uses.uses.

Module context (uses : uses_spec) <: context_spec(uses).

(* The context is implemented using maps.
    This *should* be good for th running time, as it's n ln n.
    Unfortunately, stdpp's gmaps can only be run with vm_compute,
    which seems to have performance problems.
*)
Module context_map_impl (uses : uses_spec) <: context_spec(uses).

End context_map_impl.
(* An implementation designed to be evaluable via simpl. *)
Module context_simpl_impl (uses : uses_spec) <: context_spec(uses).
    Inductive contextT {Σ} : Type :=
        | Empty
        | Var of string & iProp Σ
        | And of contextT & contextT
        | Sep of contextT & contextT
        | Constrain of Prop & contextT
    .
    Definition context Σ : Type := contextT (Σ := Σ).

End context.

    Fixpoint prop {Σ} (c : context Σ) : iProp Σ := match c with
        | Empty => True
        | Var _ T => T
        | And c1 c2 => prop c1 ∧ prop c2
        | Sep c1 c2 => prop c1 ∗ prop c2
        | Constrain P c => ∃ _ : P, prop c
    end.
    Definition var {Σ} : string -> iProp Σ -> context Σ := Var.
    Theorem var_spec {Σ} v (T : iProp Σ) : prop (var v T) ⊢ T.
    Proof. done. Qed.
    Definition constrain {Σ} (P : Prop) (ctx : context Σ) : context Σ := 
        if ctx is Constrain Q c then Constrain (P ∧ Q) c else Constrain P ctx.
    Theorem constrain_spec {Σ} P (ctx : context Σ)
        : prop (constrain P ctx) ⊢ ∃ _ : P, prop ctx.
    Proof. case ctx; intros; try done; simpl.
        apply: bi.exist_elim. move=> [w1 w2].
        by do 2 apply: bi.exist_intro'.
    Qed.
    Definition sep {Σ} (c1 c2 : context Σ) := 
        if c1 is Empty then c2 else if c2 is Empty then c1 else Sep c1 c2.
    
    Theorem sep_spec {Σ} (c1 c2 : context Σ)
        : prop (sep c1 c2) ⊢ prop c1 ∗ prop c2.
    Proof.
        by case c1; case c2; intros; simpl; 
            try rewrite bi.emp_sep; try rewrite bi.sep_emp.
    Qed.
    Definition and {Σ} (c1 c2 : context Σ) := 
        if c1 is Empty then c2 else if c2 is Empty then c1 else And c1 c2.
    Theorem and_spec {Σ} (c1 c2 : context Σ) 
        : prop (and c1 c2) ⊢ prop c1 ∧ prop c2.
    Proof.
        by case c1; case c2; intros; simpl;
            try rewrite bi.True_and; try rewrite bi.and_True.
    Qed.
    Fixpoint bind_help {Σ} (var : string) (c : context Σ) : context Σ 
    := match c with
        | Empty => Empty
        | Var var' T => if decide (var = var') then Empty else Var var' T
        | And c1 c2 => and (bind_help var c1) (bind_help var c2)
        | Sep c1 c2 => sep (bind_help var c1) (bind_help var c2)
        | Constrain P c => Constrain P (bind_help var c)
    end.
    Fixpoint lookup_help {Σ} (var : string) (c : context Σ)
        : option (iProp Σ * uses.uses * Prop)
        := match c with
            | Empty => None
            | Var var' P =>
                if decide (var = var') 
                then Some (P, uses.one, True)
                else None
            | And c1 c2 => union_with
                (λ '(P1, u1, constraints1) '(P2, u2, constraints2), Some
                    ( P1, uses.and u1 u2, constraints1 ∧ constraints2 ∧ P1 = P2)
                )
                (lookup_help var c1)
                (lookup_help var c2)
            | Sep c1 c2 => union_with
                (λ '(P1, u1, constraints1) '(P2, u2, constraints2), Some
                    ( P1, uses.sep u1 u2, constraints1 ∧ constraints2 ∧ P1 = P2)
                )
                (lookup_help var c1)
                (lookup_help var c2)
            | Constrain P c => lookup_help var c
        end.
    Lemma bind_help_spec {Σ} var (c : context Σ)
        : match lookup_help var c with
            | None => prop (bind_help var c) ⊢ prop c
            | Some (P, uses, constraints) => 
                constraints -> prop (bind_help var c) ∗ uses.prop uses P ⊢ prop c
        end.
    Proof.
        elim c; simpl.
        - done.
        - move=> var' P.
            case: (decide (var = var')).
            + by rewrite uses.one_spec bi.emp_sep.
            + done.
        - move=> c1; case: (lookup_help var c1) => [[[u1 P1] constraints1]|] H1;
            move=> c2; case: (lookup_help var c2) => [[[u2 P2] constraints2]|] H2;
            simpl; rewrite and_spec; simpl;
            [ move=>[/H1 w1 [/H2 w2 eq]] 
            | move: H2 => w2 /H1 w1 
            | move: H1 => w1 /H2 w2 
            | move: H1 H2 => w1 w2
            ];
            (apply: transitivity; last (apply: bi.and_mono; [apply: w1 | apply: w2]));
            last done; (apply: transitivity; first apply: bi.sep_and_r).
            + rewrite eq.
                apply: bi.and_mono; apply: bi.sep_mono_r; rewrite uses.and_spec;
                [apply: bi.and_elim_l | apply: bi.and_elim_r].
            + apply: bi.and_mono_r.
                rewrite -{2}[prop _]bi.sep_emp.
                apply: bi.sep_mono_r.
                apply: bi.True_intro.
            + apply: bi.and_mono_l.
                rewrite -{2}[prop _]bi.sep_emp.
                apply: bi.sep_mono_r.
                apply: bi.True_intro.
        - move=> c1; case: (lookup_help var c1) => [[[u1 P1] constraints1]|] H1;
            move=> c2; case: (lookup_help var c2) => [[[u2 P2] constraints2]|] H2;
            simpl; rewrite sep_spec; simpl;
            [ move=>[/H1 w1 [/H2 w2 eq]] 
            | move: H2 => w2 /H1 w1 
            | move: H1 => w1 /H2 w2 
            | move: H1 H2 => w1 w2
            ];
            (apply: transitivity; last (apply: bi.sep_mono; [apply: w1 | apply: w2])).
            + rewrite eq uses.sep_spec.
                rewrite bi.sep_assoc bi.sep_assoc; apply: bi.sep_mono_l.
                rewrite -bi.sep_assoc -bi.sep_assoc; apply: bi.sep_mono_r.
                by rewrite bi.sep_comm.
            + by rewrite 
                -bi.sep_assoc 
                [(prop (bind_help var c2) ∗ _)%I]bi.sep_comm
                bi.sep_assoc.
            + by rewrite bi.sep_assoc.
            + done.
        - move=> P c1.
            case: (lookup_help var c1) => [[[u T] constraints]|].
            + move=> H; move /H; clear H.
                move=> H.
                rewrite bi.sep_exist_r.
                apply: bi.exist_elim => w.
                by apply: bi.exist_intro'.
            + move=> H.
                apply: bi.exist_elim => w.
                by apply: bi.exist_intro'.
    Qed.
    Definition bind {Σ} (var : string) (ctx : context Σ) : context Σ :=
        let c := bind_help var ctx in
        match lookup_help var ctx with
            | Some (_, _, constraints) => constrain constraints c
            | None => c
        end.
    Definition lookup {Σ} (var : string) (ctx : context Σ)
        : option (iProp Σ * uses.uses)
    := option_map fst (lookup_help var ctx).
    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 /lookup /bind.
        move: (bind_help_spec var c).
        case: (lookup_help var c) => [[[T1 u1] constraints1]|] H; simpl.
        - apply: transitivity; first (apply: bi.sep_mono_r; apply: constrain_spec).
            rewrite bi.sep_exist_l.
            apply: bi.exist_elim.
            by rewrite bi.sep_comm.
        - by rewrite bi.emp_sep.
    Qed.
    Definition toplevel {Σ} (ctx : context Σ) : Prop := match ctx with
        | Empty => True
        | Constrain P Empty => P
        | _ => error ("Variables not in scope.")
    end.
    Definition toplevel_spec {Σ} (ctx : context Σ) : toplevel ctx -> ⊢ prop ctx
        := match ctx with
            | Empty => fun 'I => bi.True_intro emp
            | Constrain P Empty => bi.exist_intro
            | _ => error_elim
        end.
End context_simpl_impl.

(* Example Proof. *)

    by vm_compute.

    rewrite /check_program; simpl.
    done.

Goal ∀ {Σ} (P Q : iProp Σ), ⊢ P -∗ Q -∗ Q ∗ P.
    move=> Σ P Q.
    apply: (check_program_spec (Lam "x" (Lam "y" (Pair (Var "y") (Var "x"))))).
    rewrite /check_program; simpl.
    done.
Qed.

- Fix the performance issue.