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.




(* 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).

    (* 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_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 (Σ := Σ).


    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.