(* Stuff that should be upstreamed to std++ *)
From stdpp Require Export gmap stringmap ssreflect.
From stdpp Require Import sorting.
From stdpp Require Import options.
Set Default Proof Using "Type*".

(* Succeeds if [t] is syntactically a constructor applied to some arguments.
Note that Coq's [is_constructor] succeeds on [S], but fails on [S n]. *)
Ltac is_app_constructor t :=
  lazymatch t with
  | ?t _ => is_app_constructor t
  | _ => is_constructor t
  end.

Lemma xorb_True b1 b2 : xorb b1 b2 ↔ ¬(b1 ↔ b2).
Proof. destruct b1, b2; naive_solver. Qed.

Definition option_to_eq_Some {A} (mx : option A) : option { x | mx = Some x } :=
  match mx with
  | Some x => Some (x ↾ eq_refl)
  | None => None
  end.

(* Premise can probably be weakened to something with [ProofIrrel]. *)
Lemma option_to_eq_Some_Some `{!EqDecision A} (mx : option A) x (H : mx = Some x) :
  option_to_eq_Some mx = Some (x ↾ H).
Proof.
  destruct mx as [x'|]; simplify_eq/=; f_equal/=.
  assert (x' = x) as Hx by congruence. destruct Hx.
  f_equal. apply (proof_irrel _).
Qed.

Definition from_sum {A B C} (f : A → C) (g : B → C) (xy : A + B) : C :=
  match xy with inl x => f x | inr y => g y end.

Global Instance maybe_String : Maybe2 String := λ s,
  if s is String a s then Some (a,s) else None.

Global Instance String_inj a : Inj (=) (=) (String a).
Proof. by injection 1. Qed.

Global Instance full_relation_dec {A} : RelDecision (λ _ _ : A, True).
Proof. unfold RelDecision. apply _. Defined.

Global Instance prod_relation_dec `{RA : relation A, RB : relation B} :
  RelDecision RA → RelDecision RB → RelDecision (prod_relation RA RB).
Proof. unfold RelDecision. apply _. Defined.

Global Hint Extern 0 (from_option _ _ _) => progress simpl : core.

Definition map_sum_with `{MapFold K A M} (f : A → nat) : M → nat :=
  map_fold (λ _, plus ∘ f) 0.
Lemma map_sum_with_lookup_le `{FinMap K M} {A} (f : A → nat) (m : M A) i x :
  m !! i = Some x → f x ≤ map_sum_with f m.
Proof.
  intros. rewrite /map_sum_with (map_fold_delete_L _ _ i x m) /=; auto with lia.
Qed.

Lemma map_Forall2_dom `{FinMapDom K M C} {A B} (P : K → A → B → Prop)
    (m1 : M A) (m2 : M B) :
  map_Forall2 P m1 m2 → dom m1 ≡ dom m2.
Proof.
  revert m2. induction m1 as [|i x1 m1 ? IH] using map_ind; intros m2.
  { intros ->%map_Forall2_empty_inv_l. by rewrite !dom_empty. }
  intros (x2 & m2' & -> & ? & ? & ?)%map_Forall2_insert_inv_l; last done.
  rewrite !dom_insert IH //.
Qed.
Lemma map_Forall2_dom_L `{FinMapDom K M C, !LeibnizEquiv C} {A B}
    (P : K → A → B → Prop) (m1 : M A) (m2 : M B) :
  map_Forall2 P m1 m2 → dom m1 = dom m2.
Proof. unfold_leibniz. apply map_Forall2_dom. Qed.

Definition map_mapM
    `{!∀ A, MapFold K A (M A), !∀ A, Empty (M A), !∀ A, Insert K A (M A)}
    `{MBind F, MRet F} {A B} (f : A → F B) (m : M A) : F (M B) :=
  map_fold (λ i x mm, y ← f x; m ← mm; mret $ <[i:=y]> m) (mret ∅) m.

Section fin_map.
  Context `{FinMap K M}.

  Lemma map_insert_inv_eq {A} {m1 m2 : M A} x v u :
    m1 !! x = None →
    m2 !! x = None →
    <[x:=v]> m1 = <[x:=u]> m2 →
    v = u ∧ m1 = m2.
  Proof.
    intros Hm1 Hm2 Heq. split.
    - assert (Huv : <[x:=v]> m1 !! x = Some v). { apply lookup_insert. }
      rewrite Heq lookup_insert in Huv. by injection Huv as ->.
    - apply map_eq. intros i.
      replace m1 with (delete x (<[x:=v]> m1)) by (apply delete_insert; done).
      replace m2 with (delete x (<[x:=u]> m2)) by (apply delete_insert; done).
      by rewrite Heq.
  Qed.
  
  Lemma map_insert_inv_ne {A} {m1 m2 : M A} x1 x2 v1 v2 :
    x1 ≠ x2 →
    m1 !! x1 = None →
    m2 !! x2 = None →
    <[x1:=v1]> m1 = <[x2:=v2]> m2 →
    m1 !! x2 = Some v2 ∧ m2 !! x1 = Some v1 ∧ delete x2 m1 = delete x1 m2.
  Proof.
    intros Hx1x2 Hm1 Hm2 Hm1m2. rewrite map_eq_iff in Hm1m2. split_and!.
    - rewrite -(lookup_insert_ne _ x1 _ v1) //  Hm1m2 lookup_insert //.
    - rewrite -(lookup_insert_ne _ x2 _ v2) // -Hm1m2 lookup_insert //.
    - apply map_eq. intros y. destruct (decide (y = x1)) as [->|];
        first rewrite lookup_delete_ne // lookup_delete //.
      destruct (decide (y = x2)) as [->|];
        first rewrite lookup_delete lookup_delete_ne //.
      rewrite !lookup_delete_ne //
        -(lookup_insert_ne m2 x2 _ v2) //
        -(lookup_insert_ne m1 x1 _ v1) //.
  Qed.

  Lemma map_mapM_empty `{MBind F, MRet F} {A B} (f : A → F B) :
    map_mapM f (∅ : M A) =@{F (M B)} mret ∅.
  Proof. unfold map_mapM. by rewrite map_fold_empty. Qed.

  Lemma map_mapM_insert `{MBind F, MRet F} {A B} (f : A → F B) (m : M A) i x :
    m !! i = None → map_first_key (<[i:=x]> m) i →
    map_mapM f (<[i:=x]> m) = y ← f x; m ← map_mapM f m; mret $ <[i:=y]> m.
  Proof. intros. rewrite /map_mapM map_fold_insert_first_key //. Qed.

  Lemma map_mapM_insert_option {A B} (f : A → option B) (m : M A) i x :
    m !! i = None →
    map_mapM f (<[i:=x]> m) = y ← f x; m ← map_mapM f m; mret $ <[i:=y]> m.
  Proof.
    intros. apply: map_fold_insert; [|done].
    intros ?? z1 z2 my ???. destruct (f z1), (f z2), my; f_equal/=.
    by apply insert_commute.
  Qed.
End fin_map.

Definition map_minimal_key `{MapFold K A M} (R : relation K) `{!RelDecision R}
    (m : M) : option K :=
  map_fold (λ i _ mj,
    match mj with
    | Some j => if decide (R i j) then Some i else Some j
    | None => Some i
    end) None m.

Section map_sorted.
  Context `{FinMap K M} (R : relation K) .

  Lemma map_minimal_key_None {A} `{!RelDecision R} (m : M A) :
    map_minimal_key R m = None ↔ m = ∅.
  Proof.
    split; [|intros ->; apply map_fold_empty].
    induction m as [|j x m ?? _] using map_first_key_ind; intros Hm; [done|].
    rewrite /map_minimal_key map_fold_insert_first_key // in Hm.
    repeat case_match; simplify_option_eq.
  Qed.

  Lemma map_minimal_key_Some_1 {A} `{!RelDecision R, !PreOrder R, !Total R}
      (m : M A) i :
    map_minimal_key R m = Some i →
      is_Some (m !! i) ∧ ∀ j, is_Some (m !! j) → R i j.
  Proof.
    revert i. induction m as [|j x m ?? IH] using map_first_key_ind; intros i Hm.
    { by rewrite /map_minimal_key map_fold_empty in Hm. }
    rewrite /map_minimal_key map_fold_insert_first_key // in Hm.
    destruct (map_fold _ _ m) as [i'|] eqn:Hfold; simplify_eq.
    - apply IH in Hfold as [??]. rewrite lookup_insert_is_Some.
      case_decide as HR; simplify_eq/=.
      + split; [by auto|]. intros j [->|[??]]%lookup_insert_is_Some; [done|].
        trans i'; eauto.
      + split.
        { right; split; [|done]. intros ->. by destruct HR. }
        intros j' [->|[??]]%lookup_insert_is_Some; [|by eauto].
        by destruct (total R i j').
    - apply map_minimal_key_None in Hfold as ->.
      split; [rewrite lookup_insert; by eauto|].
      intros j' [->|[? Hj']]%lookup_insert_is_Some; [done|].
      rewrite lookup_empty in Hj'. by destruct Hj'.
  Qed.

  Lemma map_minimal_key_Some {A} `{!RelDecision R, !PartialOrder R, !Total R}
      (m : M A) i :
    map_minimal_key R m = Some i ↔
      is_Some (m !! i) ∧ ∀ j, is_Some (m !! j) → R i j.
  Proof.
    split; [apply map_minimal_key_Some_1|].
    intros [Hi ?]. destruct (map_minimal_key R m) as [i'|] eqn:Hmin.
    - f_equal. apply map_minimal_key_Some_1 in Hmin as [??].
      apply (anti_symm R); eauto.
    - apply map_minimal_key_None in Hmin as ->.
      rewrite lookup_empty in Hi. by destruct Hi.
  Qed.

  Lemma map_sorted_ind {A} `{!PreOrder R, !Total R} (P : M A → Prop) :
    P ∅ →
    (∀ i x m,
      m !! i = None →
      (∀ j, is_Some (m !! j) → R i j) →
      P m →
      P (<[i:=x]> m)) →
    (∀ m, P m).
  Proof.
    intros Hemp Hins m. induction (Nat.lt_wf_0_projected size m) as [m _ IH].
    cut (m = ∅ ∨ map_Exists (λ i _, ∀ j, is_Some (m !! j) → R i j) m).
    { intros [->|(i & x & Hi & ?)]; [done|]. rewrite -(insert_delete m i x) //.
      apply Hins; [by rewrite lookup_delete|..].
      - intros j ?%lookup_delete_is_Some. naive_solver.
      - apply IH.
        rewrite -{2}(insert_delete m i x) // map_size_insert lookup_delete. lia. }
    clear P Hemp Hins IH. induction m as [|i x m ? IH] using map_ind; [by auto|].
    right. destruct IH as [->|(i' & x' & ? & ?)].
    { rewrite insert_empty map_Exists_singleton.
      by intros j [y [-> ->]%lookup_singleton_Some]. }
    apply map_Exists_insert; first done. destruct (total R i i').
    - left. intros j [->|[??]]%lookup_insert_is_Some; [done|]. trans i'; eauto.
    - right. exists i', x'. split; [done|].
      intros j [->|[??]]%lookup_insert_is_Some; eauto.
  Qed.
End map_sorted.

Definition map_fold_sorted `{!MapFold K A M} {B}
    (R : relation K) `{!RelDecision R}
    (f : K → A → B → B) (b : B)
    (m : M) : B := foldr (λ '(i,x), f i x) b $
  merge_sort (prod_relation R (λ _ _, True)) (map_to_list m).

Definition map_mapM_sorted
    `{!∀ A, MapFold K A (M A), !∀ A, Empty (M A), !∀ A, Insert K A (M A)}
    `{MBind F, MRet F} {A B}
    (R : relation K) `{!RelDecision R}
    (f : A → F B) (m : M A) : F (M B) :=
  map_fold_sorted R (λ i x mm, y ← f x; m ← mm; mret $ <[i:=y]> m) (mret ∅) m.

Section fin_map.
  Context `{FinMap K M}.
  Context (R : relation K) `{!RelDecision R, !PartialOrder R, !Total R}.

  Lemma map_fold_sorted_empty {A B} (f : K → A → B → B) b :
    map_fold_sorted R f b (∅ : M A) = b.
  Proof. by rewrite /map_fold_sorted map_to_list_empty. Qed.

  Lemma map_fold_sorted_insert {A B} (f : K → A → B → B) (m : M A) b i x :
    m !! i = None → (∀ j, is_Some (m !! j) → R i j) →
    map_fold_sorted R f b (<[i:=x]> m) = f i x (map_fold_sorted R f b m).
  Proof.
    intros Hi Hleast. unfold map_fold_sorted.
    set (R' := prod_relation R _).
    assert (PreOrder R').
    { split; [done|].
      intros [??] [??] [??] [??] [??]; split; [by etrans|done]. }
    assert (Total R').
    { intros [i1 ?] [i2 ?]. destruct (total R i1 i2); [by left|by right]. }
    assert (merge_sort R' (map_to_list (<[i:=x]> m))
      = (i,x) :: merge_sort R' (map_to_list m)) as ->; [|done].
    eapply (Sorted_unique_strong R').
    - intros [i1 y1] [i2 y2].
      rewrite !merge_sort_Permutation elem_of_cons !elem_of_map_to_list.
      rewrite lookup_insert_Some. intros ?? [? _] [? _].
      assert (i1 = i2) as -> by (by apply (anti_symm R)); naive_solver.
    - apply (Sorted_merge_sort _).
    - apply Sorted_cons; [apply (Sorted_merge_sort _)|].
      destruct (merge_sort R' (map_to_list m))
        as [|[i' x'] ixs] eqn:Hixs; repeat constructor; simpl.
      apply Hleast. exists x'. apply elem_of_map_to_list.
      rewrite -(merge_sort_Permutation R' (map_to_list m)) Hixs. left.
    - by rewrite !merge_sort_Permutation map_to_list_insert.
  Qed.

  Lemma map_mapM_sorted_empty `{MBind F, MRet F} {A B} (f : A → F B) :
    map_mapM_sorted R f (∅ : M A) =@{F (M B)} mret ∅.
  Proof. by rewrite /map_mapM_sorted map_fold_sorted_empty. Qed.

  Lemma map_mapM_sorted_insert `{MBind F, MRet F}
      {A B} (f : A → F B) (m : M A) i x :
    m !! i = None → (∀ j, is_Some (m !! j) → R i j) →
    map_mapM_sorted R f (<[i:=x]> m)
    = y ← f x; m ← map_mapM_sorted R f m; mret $ <[i:=y]> m.
  Proof. intros. by rewrite /map_mapM_sorted map_fold_sorted_insert. Qed.
End fin_map.