Require Import Coq.Structures.Equalities.

Require Import TVSMT.Maps.
Require Import TVSMT.Errors.
Require Import TVSMT.Common.

#[local] Open Scope positive.

#[local] Hint Resolve in_eq : core.
#[local] Hint Resolve in_cons : core.

Declare Scope monad_scope.

Module Type Monad.

  Parameter mon : Type -> Type.

  Parameter ret : forall (A : Type) (x : A), mon A.
  Arguments ret [A].

  Parameter bind : forall (A B : Type) (f : mon A) (g : A -> mon B), mon B.
  Arguments bind [A B].

  Parameter bind2 : forall (A B C: Type) (f: mon (A * B)) (g: A -> B -> mon C), mon C.
  Arguments bind2 [A B C].

End Monad.

Module MonadExtra(M : Monad).
  Import M.

  Module Import MonadNotation.

    Notation "'do' X <- A ; B" :=
      (bind A (fun X => B))
        (at level 200, X ident, A at level 100, B at level 200) : monad_scope.
    Notation "'do' ( X , Y ) <- A ; B" :=
      (bind2 A (fun X Y => B))
        (at level 200, X ident, Y ident, A at level 100, B at level 200) : monad_scope.

  End MonadNotation.

  #[local] Open Scope monad_scope.

  Fixpoint traverselist {A B: Type} (f: A -> mon B) (l: list A) {struct l}: mon (list B) :=
    match l with
    | nil => ret nil
    | x::xs =>
      do r <- f x;
      do rs <- traverselist f xs;
      ret (r::rs)
    end.

  Fixpoint collectlist {A : Type} (f : A -> mon unit) (l : list A) {struct l} : mon unit :=
    match l with
    | nil => ret tt
    | x::xs => do y <- f x; collectlist f xs
    end.

  Definition mfold_left {A B} (f: A -> B -> mon A) (l: list B) (s: mon A): mon A :=
    fold_left (fun a b => do a' <- a; f a' b) l s.

End MonadExtra.

Module Type State.
  Parameter st : Type.
  Parameter st_prop : st -> st -> Prop.

  Axiom st_refl : forall s, st_prop s s.
  Axiom st_trans : forall s1 s2 s3, st_prop s1 s2 -> st_prop s2 s3 -> st_prop s1 s3.
End State.

Module Statemonad(S : State) <: Monad.

  Inductive res (A: Type) (s: S.st): Type :=
  | Error : Errors.errmsg -> res A s
  | OK : A -> forall (s' : S.st), S.st_prop s s' -> res A s.

  Arguments OK [A s].
  Arguments Error [A s].

  Definition mon (A: Type) : Type := forall (s: S.st), res A s.

  Definition ret {A: Type} (x: A) : mon A :=
    fun (s : S.st) => OK x s (S.st_refl s).

  Definition bind {A B: Type} (f: mon A) (g: A -> mon B) : mon B :=
    fun (s : S.st) =>
      match f s with
      | Error msg => Error msg
      | OK a s' i =>
        match g a s' with
        | Error msg => Error msg
        | OK b s'' i' => OK b s'' (S.st_trans s s' s'' i i')
        end
      end.

  Definition bind2 {A B C: Type} (f: mon (A * B)) (g: A -> B -> mon C) : mon C :=
    bind f (fun xy => g (fst xy) (snd xy)).

  Definition handle_error {A: Type} (f g: mon A) : mon A :=
    fun (s : S.st) =>
      match f s with
      | OK a s' i => OK a s' i
      | Error _ => g s
      end.

  Definition error {A: Type} (err: Errors.errmsg) : mon A := fun (s: S.st) => Error err.

  Definition get : mon S.st := fun s => OK s s (S.st_refl s).

  Definition set (s: S.st) (i: forall s', S.st_prop  s' s) : mon unit :=
    fun s' => OK tt s (i s').

  Definition run_mon {A: Type} (s: S.st) (m: mon A): Errors.res A :=
    match m s with
    | OK a s' i => Errors.OK a
    | Error err => Errors.Error err
    end.

End Statemonad.

Definition max_key {A} (t: PTree.t A) :=
  fold_right Pos.max 1%positive (map fst (PTree.elements t)).

Lemma max_key_correct' :
  forall l hi, In hi l -> hi <= fold_right Pos.max 1 l.
Proof.
  induction l; crush.
  inv H. lia.
  destruct (Pos.max_dec a (fold_right Pos.max 1 l)); rewrite e.
  - apply Pos.max_l_iff in e.
    assert (forall a b c, a <= c -> c <= b -> a <= b) by lia.
    eapply H; eauto.
  - apply IHl; auto.
Qed.

Lemma max_key_correct :
  forall A h_tree hi (c: A),
    h_tree ! hi = Some c ->
    hi <= max_key h_tree.
Proof.
  unfold max_key. intros. apply PTree.elements_correct in H.
  apply max_key_correct'.
  eapply in_map with (f := fst) in H. auto.
Qed.

Lemma max_not_present :
  forall A k (h: PTree.t A), k > max_key h -> h ! k = None.
Proof.
  intros. destruct (h ! k) eqn:?; auto.
  apply max_key_correct in Heqo. lia.
Qed.

Lemma filter_none :
  forall A f l (x: A), filter f l = nil -> In x l -> f x = false.
Proof. induction l; crush; inv H0; subst; destruct_match; crush. Qed.

Lemma filter_set :
  forall A l l' f (x: A),
    (In x l -> In x l') ->
    In x (filter f l) ->
    In x (filter f l').
Proof.
  induction l; crush.
  destruct_match; crush. inv H0; crush.
  apply filter_In. simplify; crush.
Qed.

Lemma filter_cons_true :
  forall A f l (a: A) l',
    filter f l = a :: l' -> f a = true.
Proof.
  induction l; crush. destruct (f a) eqn:?.
  inv H. auto. eapply IHl; eauto.
Qed.

Lemma PTree_set_elements :
  forall A t x x' (c: A),
    In x (PTree.elements t) ->
    x' <> (fst x) ->
    In x (PTree.elements (PTree.set x' c t)).
Proof.
  intros. destruct x.
  eapply PTree.elements_correct. simplify.
  rewrite PTree.gso; auto. apply PTree.elements_complete in H. auto.
Qed.

Lemma filter_set2 :
  forall A x y z (h: PTree.t A),
    In z (PTree.elements (PTree.set x y h)) ->
    In z (PTree.elements h) \/ fst z = x.
Proof.
  intros. destruct z.
  destruct (Pos.eq_dec p x); subst.
  tauto.
  left. apply PTree.elements_correct. apply PTree.elements_complete in H.
  rewrite PTree.gso in H; auto.
Qed.

Lemma in_filter : forall A f l (x: A), In x (filter f l) -> In x l.
Proof. induction l; crush. destruct_match; crush. inv H; crush. Qed.

Lemma filter_norepet:
  forall A f (l: list A),
    list_norepet l ->
    list_norepet (filter f l).
Proof.
  induction l; crush.
  inv H. destruct (f a).
  constructor. unfold not in *; intros. apply H2.
  eapply in_filter; eauto.
  apply IHl; auto.
  apply IHl; auto.
Qed.

Lemma filter_norepet2:
  forall A B g (l: list (A * B)),
    list_norepet (map fst l) ->
    list_norepet (map fst (filter g l)).
Proof.
  induction l; crush.
  inv H. destruct (g a) eqn:?.
  simplify. constructor. unfold not in *. intros.
  eapply H2.
  apply list_in_map_inv in H. simplify; subst.
  rewrite H.
  apply filter_In in H1. simplify.
  apply in_map. eauto.
  eapply IHl. eauto.
  eapply IHl. eauto.
Qed.

Module Type Hashable := UsualDecidableType.

Module HashTree(Import H: Hashable).

  Definition hash := positive.
  Definition hash_tree := PTree.t t.

  Definition empty := PTree.empty t.

  Definition find_tree (el: t) (h: hash_tree) : option hash :=
    match filter (fun x => if eq_dec el (snd x) then true else false) (PTree.elements h) with
    | (p, _) :: nil => Some p
    | _ => None
    end.

  Definition hash_value (max: hash) (e: t) (h: hash_tree): hash * hash_tree :=
    match find_tree e h with
    | Some p => (p, h)
    | None =>
      let nkey := Pos.max max (max_key h) + 1 in
      (nkey, PTree.set nkey e h)
    end.

  Definition hash_value2 (max: hash) (e: t * t) (h: hash_tree): (hash * hash) * hash_tree :=
    let (e1, e2) := e in
    let (v1, h1) := hash_value max e1 h in
    let (v2, h2) := hash_value max e2 h1 in
    ((v1, v2), h2).

  Definition hash_value3 (max: hash) (e: t * t * t) (h: hash_tree): (hash * hash * hash) * hash_tree :=
    let '(e1, e2, e3) := e in
    let (v1, h1) := hash_value max e1 h in
    let (v2, h2) := hash_value max e2 h1 in
    let (v3, h3) := hash_value max e3 h2 in
    ((v1, v2, v3), h3).

  Definition hash_value4 (max: hash) (e: t * t * t * t) (h: hash_tree): (hash * hash * hash * hash) * hash_tree :=
    let '(e1, e2, e3, e4) := e in
    let (v1, h1) := hash_value max e1 h in
    let (v2, h2) := hash_value max e2 h1 in
    let (v3, h3) := hash_value max e3 h2 in
    let (v4, h4) := hash_value max e4 h3 in
    ((v1, v2, v3, v4), h4).

  Definition hash_list (max: hash) (e: list t) (h: hash_tree): list hash * hash_tree :=
    fold_left (fun s e =>
                 let (i, h') := hash_value max e (snd s) in
                 (fst s ++ i::nil, h')
              ) e (nil, h).

  Definition wf_hash_table h_tree :=
    (forall x c, h_tree ! x = Some c -> find_tree c h_tree = Some x).

  Lemma find_tree_correct :
    forall c h_tree p,
      find_tree c h_tree = Some p ->
      h_tree ! p = Some c.
  Proof.
    intros.
    unfold find_tree in H. destruct_match; crush.
    destruct_match; simplify.
    destruct_match; crush.
    assert (In (p, t0) (filter
                          (fun x : hash * t =>
                             if eq_dec c (snd x) then true else false) (PTree.elements h_tree))).
    { setoid_rewrite Heql. constructor; auto. }
    apply filter_In in H. simplify. destruct_match; crush. subst.
    apply PTree.elements_complete; auto.
  Qed.

  Lemma find_tree_unique :
    forall c h_tree p p',
      find_tree c h_tree = Some p ->
      h_tree ! p' = Some c ->
      p = p'.
  Proof.
    intros.
    unfold find_tree in H.
    repeat (destruct_match; crush; []).
    assert (In (p, t0) (filter
                          (fun x : hash * t =>
                             if eq_dec c (snd x) then true else false) (PTree.elements h_tree))).
    { setoid_rewrite Heql. constructor; auto. }
    apply filter_In in H. simplify.
    destruct (Pos.eq_dec p p'); auto.
    exfalso.
    destruct_match; subst; crush.
    assert (In (p', t0) (PTree.elements h_tree) /\ (fun x : hash * t =>
                                                      if eq_dec t0 (snd x) then true else false) (p', t0) = true).
    { split. apply PTree.elements_correct. auto. setoid_rewrite Heqs. auto. }
    apply filter_In in H. setoid_rewrite Heql in H. inv H. simplify. crush. crush.
  Qed.

  Lemma hash_no_element' :
    forall c h_tree,
      find_tree c h_tree = None ->
      wf_hash_table h_tree ->
      ~ forall x, h_tree ! x = Some c.
  Proof.
    unfold not, wf_hash_table; intros.
    specialize (H1 1). eapply H0 in H1. crush.
  Qed.

  Lemma hash_no_element :
    forall c h_tree,
      find_tree c h_tree = None ->
      wf_hash_table h_tree ->
      ~ exists x, h_tree ! x = Some c.
  Proof.
    unfold not, wf_hash_table; intros.
    simplify. apply H0 in H2. rewrite H in H2. crush.
  Qed.

  Lemma wf_hash_table_set_gso' :
    forall h x p0 c',
      filter
        (fun x : hash * t =>
           if eq_dec c' (snd x) then true else false) (PTree.elements h) =
      (x, p0) :: nil ->
      h ! x = Some p0 /\ p0 = c'.
  Proof.
    intros.
    match goal with
    | H: filter ?f ?el = ?x::?xs |- _ =>
      assert (In x (filter f el)) by (rewrite H; crush)
    end.
    apply filter_In in H0. simplify. destruct_match; subst; crush.
    apply PTree.elements_complete; auto.
    destruct_match; crush.
  Qed.

  Lemma wf_hash_table_set_gso :
    forall x x' c' c h,
      x <> x' ->
      wf_hash_table h ->
      find_tree c' h = Some x ->
      find_tree c h = None ->
      find_tree c' (PTree.set x' c h) = Some x.
  Proof.
    intros. pose proof H1 as X. unfold find_tree in H1.
    destruct_match; crush.
    destruct p. destruct l; crush.
    apply wf_hash_table_set_gso' in Heql. simplify.
    pose proof H2 as Z. apply hash_no_element in H2; auto.
    destruct (eq_dec c c'); subst.
    { exfalso. eapply H2. econstructor; eauto. }
    unfold wf_hash_table in H0.
    assert (In (x', c) (PTree.elements (PTree.set x' c h))).
    { apply PTree.elements_correct. rewrite PTree.gss; auto. }
    assert (In (x, c') (PTree.elements h)).
    { apply PTree.elements_correct; auto. }
    assert (In (x, c') (PTree.elements (PTree.set x' c h))).
    { apply PTree.elements_correct. rewrite PTree.gso; auto. }
    pose proof X as Y.
    unfold find_tree in X. repeat (destruct_match; crush; []).
    match goal with
    | H: filter ?f ?el = ?x::?xs |- _ =>
      assert (In x (filter f el)) by (rewrite H; crush)
    end.
    apply filter_In in H6. simplify. destruct_match; crush; subst.
    unfold find_tree. repeat (destruct_match; crush).
    { eapply filter_none in Heql0.
      2: { apply PTree.elements_correct. rewrite PTree.gso; eauto. }
      destruct_match; crush. }

    { subst.
      repeat match goal with
             | H: filter ?f ?el = ?x::?xs |- _ =>
               learn H; assert (In x (filter f el)) by (rewrite H; crush)
             end.
      eapply filter_set in H10. rewrite Heql0 in H10. inv H10. simplify. auto.
      inv H11. auto. }

    { exfalso. subst.
      repeat match goal with
             | H: filter ?f ?el = ?x::?xs |- _ =>
               learn H; assert (In x (filter f el)) by (rewrite H; crush)
             end.

      pose proof H8 as X2. destruct p1.
      pose proof X2 as X4.
      apply in_filter in X2. apply PTree.elements_complete in X2.
      assert (In (p, t2) (filter
                            (fun x : positive * t => if eq_dec t0 (snd x) then true else false)
                            (PTree.elements (PTree.set x' c h)))) by (rewrite H6; eauto).
      pose proof H11 as X3.
      apply in_filter in H11. apply PTree.elements_complete in H11.
      destruct (peq p0 p); subst.
      {
        assert (list_norepet (map fst (filter
                                         (fun x : positive * t => if eq_dec t0 (snd x) then true else false)
                                         (PTree.elements (PTree.set x' c h))))).
        { eapply filter_norepet2. eapply PTree.elements_keys_norepet. }
        rewrite Heql0 in H12. simplify. inv H12. eapply H15. apply in_eq.
      }
      { apply filter_In in X4. simplify. destruct_match; crush; subst.
        apply filter_In in X3. simplify. destruct_match; crush; subst.
        destruct (peq p x'); subst.
        { rewrite PTree.gss in H11; crush. }
        { destruct (peq p0 x'); subst.
          { rewrite PTree.gss in X2; crush. }
          { rewrite PTree.gso in X2 by auto.
            rewrite PTree.gso in H11 by auto.
            assert (p = p0) by (eapply find_tree_unique; eauto).
            crush. } } } }
  Qed.

  Lemma wf_hash_table_set :
    forall h_tree c v (GT: v > max_key h_tree),
      find_tree c h_tree = None ->
      wf_hash_table h_tree ->
      wf_hash_table (PTree.set v c h_tree).
  Proof.
    unfold wf_hash_table; simplify.
    destruct (peq v x); subst.
    pose proof (hash_no_element c h_tree H H0).
    rewrite PTree.gss in H1. simplify.
    unfold find_tree.
    assert (In (x, c0) (PTree.elements (PTree.set x c0 h_tree))
            /\ (fun x : positive * t => if eq_dec c0 (snd x) then true else false)
                 (x, c0) = true).
    { simplify. apply PTree.elements_correct. rewrite PTree.gss. auto.
      destruct (eq_dec c0 c0); crush. }
    destruct_match.
    apply filter_In in H1. rewrite Heql in H1. crush.
    apply filter_In in H1. repeat (destruct_match; crush; []). subst.
    destruct l. simplify. rewrite Heql in H1. inv H1. inv H3. auto.
    crush.

    exfalso. apply H2. destruct p.
    pose proof Heql as X. apply filter_cons_true in X. destruct_match; crush; subst.
    assert (In (p0, t0) (filter
                           (fun x : positive * t => if eq_dec t0 (snd x) then true else false)
                           (PTree.elements (PTree.set x t0 h_tree)))) by (rewrite Heql; eauto).
    assert (In (p, t1) (filter
                          (fun x : positive * t => if eq_dec t0 (snd x) then true else false)
                          (PTree.elements (PTree.set x t0 h_tree)))) by (rewrite Heql; eauto).
    apply filter_In in H4. simplify. destruct_match; crush; subst.
    apply in_filter in H3. apply PTree.elements_complete in H5. apply PTree.elements_complete in H3.
    assert (list_norepet (map fst (filter
                                     (fun x : positive * t => if eq_dec t1 (snd x) then true else false)
                                     (PTree.elements (PTree.set x t1 h_tree))))).
    { eapply filter_norepet2. eapply PTree.elements_keys_norepet. }
    rewrite Heql in H4. simplify.
    destruct (peq p0 p); subst.
    { inv H4. exfalso. eapply H8. eauto. }
    destruct (peq x p); subst.
    rewrite PTree.gso in H3; auto. econstructor; eauto.
    rewrite PTree.gso in H5; auto. econstructor; eauto.

    rewrite PTree.gso in H1; auto.
    destruct (eq_dec c c0); subst.
    { apply H0 in H1. rewrite H in H1. discriminate. }
    apply H0 in H1.
    apply wf_hash_table_set_gso; eauto.
  Qed.

  Lemma wf_hash_table_distr :
    forall m p h_tree h h_tree',
      hash_value m p h_tree = (h, h_tree') ->
      wf_hash_table h_tree ->
      wf_hash_table h_tree'.
  Proof.
    unfold hash_value; simplify.
    destruct_match.
    - inv H; auto.
    - inv H. apply wf_hash_table_set; try lia; auto.
  Qed.

  Lemma wf_hash_table_eq :
    forall h_tree a b c,
      wf_hash_table h_tree ->
      h_tree ! a = Some c ->
      h_tree ! b = Some c ->
      a = b.
  Proof.
    unfold wf_hash_table; intros; apply H in H0; eapply find_tree_unique; eauto.
  Qed.

  Lemma hash_constant :
    forall p h h_tree hi c h_tree' m,
      h_tree ! hi = Some c ->
      hash_value m p h_tree = (h, h_tree') ->
      h_tree' ! hi = Some c.
  Proof.
    intros. unfold hash_value in H0. destruct_match.
    inv H0. eauto.
    inv H0.
    pose proof H. apply max_key_correct in H0.
    rewrite PTree.gso; solve [eauto | lia].
  Qed.

  Lemma find_tree_Some :
    forall el h v,
      find_tree el h = Some v ->
      h ! v = Some el.
  Proof.
    intros. unfold find_tree in *.
    destruct_match; crush. destruct p.
    destruct_match; crush.
    match goal with
    | H: filter ?f ?el = ?x::?xs |- _ =>
      assert (In x (filter f el)) by (rewrite H; crush)
    end.
    apply PTree.elements_complete.
    apply filter_In in H. inv H.
    destruct_match; crush.
  Qed.

  Lemma hash_present_eq :
    forall m e1 e2 p1 h h',
      hash_value m e2 h = (p1, h') ->
      h ! p1 = Some e1 -> e1 = e2.
  Proof.
    intros. unfold hash_value in *. destruct_match.
    - inv H. apply find_tree_Some in Heqo.
      rewrite Heqo in H0. inv H0. auto.
    - inv H. assert (h ! (Pos.max m (max_key h) + 1) = None)
        by (apply max_not_present; lia). crush.
  Qed.

  Module HashState <: State.
    Definition st := hash_tree.
    Definition st_prop (h1 h2: hash_tree): Prop :=
      forall x y, h1 ! x = Some y -> h2 ! x = Some y.

    Lemma st_refl :
      forall s, st_prop s s.
    Proof. unfold st_prop; auto. Qed.

    Lemma st_trans :
      forall s1 s2 s3,
        st_prop s1 s2 -> st_prop s2 s3 -> st_prop s1 s3.
    Proof.
      unfold st_prop; intros; eauto.
    Qed.

    #[global] Instance HashStatePreorder : PreOrder st_prop :=
      { PreOrder_Reflexive := st_refl;
        PreOrder_Transitive := st_trans;
      }.
  End HashState.

  Module HashMonad := Statemonad(HashState).

End HashTree.

Definition gt_1 {A} h :=
  forall x (y: A), h ! x = Some y -> 1 < x.

Module HashTreeProperties (Import H: Hashable).

  Module Import HT := HashTree(H).

  Lemma hash_value_gt :
    forall max v h,
      gt_1 h ->
      gt_1 (snd (hash_value max v h)).
  Proof.
    unfold gt_1, hash_value; intros.
    destruct_match; eauto.
    destruct (peq (Pos.max max (max_key h) + 1) x); [subst;lia|].
    cbn [snd] in *. rewrite PTree.gso in H0; eauto.
  Qed.

End HashTreeProperties.