Module PositiveMapAddOn

Require Export BinNums.
Require Import PArith.BinPosDef.
Require Import PArith.BinPos.
Require Export OptionMonad.
Require Export FMapPositive.
Require Import List.
Require Import CoqAddOn.
Require Import Sorting.Permutation.

Local Open Scope type_scope.
Local Open Scope list_scope.
Import List.ListNotations.

Import PositiveMap.

Hint Unfold E.eq.
Hint Resolve f_equal f_equal2.

Implicit Arguments Leaf [A].

Fixpoint single {A} (x:positive) (v:A) : t A :=
  match x with
  | xH => Node Leaf (Some v) Leaf
  | xO x' => Node (single x' v) None Leaf
  | xI x' => Node Leaf None (single x' v)
  end.

Lemma append_assoc i: forall j k, append (append i j) k = append i (append j k).
Hint Rewrite append_assoc append_neutral_r append_neutral_l app_nil_r @SOME_assoc: posmap.

Lemma xelements_single {A} x (v:A): forall path, xelements (single x v) path = [(append path x,v)].

Theorem elements_single {A} x (v:A): elements (single x v) = [(x,v)].

Lemma single_equiv_add_Leaf {A} x (v:A): single x v = add x v Leaf.




Section SimplerMap.

  Context {A B: Type}.

  Variable f : A -> B.

  Fixpoint map (m : t A) : t B :=
       match m with
        | Leaf => Leaf
        | Node l o r => Node (map l)
                             (option_map f o)
                             (map r)
       end.

Lemma xelements_map (m: t A): forall path,
  xelements (map m) path = List.map (fun p => (fst p, f (snd p))) (xelements m path).

Theorem elements_map (m: t A):
  elements (map m) = List.map (fun p => (fst p, f (snd p))) (elements m).

End SimplerMap.


Section SimplerEqual.

Context {A: Type}.
Variable cmp: A -> A -> bool.
Hypothesis cmp_correct: forall x y, If cmp x y THEN x=y.

Fixpoint equal (m1 m2 : t A) : bool :=
    match m1, m2 with
      | Leaf, Leaf => true
      | Node l1 o1 r1, Node l2 o2 r2 =>
           (match o1, o2 with
             | None, None => true
             | Some v1, Some v2 => cmp v1 v2
             | _, _ => false
            end)
           && equal l1 l2 && equal r1 r2
       | _,_ => false
     end.

Lemma equal_correct m1: forall m2, If equal m1 m2 THEN m1=m2.

End SimplerEqual.


Lemma elements_spec {A} m x (v:A): List.In (x,v) (elements m) <-> find x m = Some v.

Lemma In_mem {A} m x (v:A): List.In (x,v) (PositiveMap.elements m) -> (PositiveMap.mem x m)=true.

Definition node {A} (l:PositiveMap.t A) o (r:PositiveMap.t A) :=
  match l,r,o with
  | Leaf,Leaf,None => Leaf
  | _ , _, _ => Node l o r
  end.

Lemma find_node {A} (l:PositiveMap.t A) o (r:PositiveMap.t A) x: find x (node l o r) = find x (Node l o r).

Lemma xelements_node {A} (l:PositiveMap.t A) o (r:PositiveMap.t A): forall path, xelements (node l o r) path = xelements (Node l o r) path.


Lemma lt_key_append_O path: forall p, E.lt (append path (xO p)) path.

Lemma lt_key_append_I path: forall p, E.lt path (append path (xI p)).

Lemma lt_key_append_O_I path: forall p1 p2, E.lt (append path (xO p1)) (append path (xI p2)).
Local Hint Resolve lt_key_append_O lt_key_append_I lt_key_append_O_I.

Lemma xelements_Forall_ltkey_O {A} (m:t A): forall path p,
  Forall (fun c => E.lt (fst c) path) (xelements m (append path (xO p))).

Lemma xelements_Forall_ltkey_I {A} (m:t A): forall path p,
  Forall (fun c => E.lt path (fst c)) (xelements m (append path (xI p))).

Lemma xelements_Forall_ltkey_O_I {A} (m:t A): forall path p1 p2,
  Forall (fun c => E.lt (append path (xO p1)) (fst c)) (xelements m (append path (xI p2))).

Lemma xelements_Forall_ltkey_I_O {A} (m:t A): forall path p1 p2,
  Forall (fun c => E.lt (fst c) (append path (xI p1))) (xelements m (append path (xO p2))).
Local Hint Resolve xelements_Forall_ltkey_O xelements_Forall_ltkey_I xelements_Forall_ltkey_O_I xelements_Forall_ltkey_I_O.

Lemma xelements_Forall_exclusion {A} (m1: t A): forall (m2:t A) path p3 p2,
  Forall (fun c => Forall (fun c' => E.lt (fst c) (fst c')) (xelements m2 (append path (xI p3)))) (xelements m1 (append path (xO p2))).
Local Hint Resolve xelements_Forall_exclusion.



Fixpoint altRemove {A} (i : positive) (m : t A) : t A :=
    match m with
      | Leaf => Leaf
      | Node l o r =>
        match i with
          | xH => node l None r
          | xO i => node (altRemove i l) o r
          | xI i => node l o (altRemove i r)
        end
    end.

Lemma altRemove_def {A} (m: t A): forall i, altRemove i m = remove i m.

Lemma append_discr_xI p1: forall p2, append p1 (xI p2) = p1 -> False.

Lemma append_discr_xO p1: forall p2, append p1 (xO p2) = p1 -> False.

Lemma append_discr_xI_xO p p1 p2: append p (xI p1) = append p (xO p2) -> False.

Definition eq_key {A} x (c:positive*A) := Pos.eq_dec x (fst c) ||| false.

Lemma Forall_lt_neq_key_1 {A} p (l: list (positive * A)):
  Forall (fun c => E.lt p (fst c)) l -> Forall (fun c => negb (eq_key p c) = true) l .
Local Hint Resolve Forall_lt_neq_key_1.

Lemma Forall_lt_neq_key_2 {A} p (l: list (positive * A)):
  Forall (fun c => E.lt (fst c) p) l -> Forall (fun c => negb (eq_key p c) = true) l.
Local Hint Resolve Forall_lt_neq_key_2.

Lemma xelements_remove {A} (m: t A) : forall x path,
  xelements (remove x m) path = List.filter (fun c => negb (eq_key (append path x) c)) (xelements m path).

Lemma elements_remove {A} x (m: t A):
  elements (remove x m) = List.filter (fun c => negb (eq_key x c)) (elements m).


Local Hint Resolve Forall_negb_true_false.

Lemma find_xelements {A} (m:t A):
 forall x path, match find x m with Some v => [ (append path x,v) ] | None => [] end =
                 List.filter (eq_key (append path x)) (xelements m path).

Lemma find_elements {A} x (m: t A) :
  match find x m with Some v => [ (x,v) ] | None => [] end = List.filter (eq_key x) (elements m).

Local Hint Resolve filter_triv_true Forall_false_negb_true filter_empty_false.

Theorem find_remove_elements_Permutation {A} x (m: t A):
  wte (find x m) (fun v => Permutation ((x,v)::(elements (remove x m))) (elements m))
                 (elements (remove x m)=elements m).


Section merge.

  Context {A : Type}.
  Variable f : A -> A -> option A.

  Definition nodeMerge (o1 o2: option A) :=
    match o1 with
    | None => o2
    | Some v1 => match o2 with
                 | None => o1
                 | Some v2 => f v1 v2
                 end
    end.

  Lemma nodeMerge_None_r (o: option A): nodeMerge o None = o.

  Fixpoint merge (m1 : t A)(m2 : t A) : t A :=
    match m1 with
    | Leaf => m2
    | Node l1 o1 r1 =>
        match m2 with
        | Leaf => m1
        | Node l2 o2 r2 => node (merge l1 l2) (nodeMerge o1 o2) (merge r1 r2)
        end
    end.

    Lemma find_merge: forall (m1 m2:t A) (i: positive),
      find i (merge m1 m2) = nodeMerge (find i m1) (find i m2).

  Inductive inMerge: list (positive*A) -> list (positive*A) -> list (positive*A) -> Prop :=
  | inM_nil_l l: inMerge [] l l
  | inM_nir_r l: inMerge l [] l
  | inM_LT c1 l1 c2 l2 l: E.lt (fst c1) (fst c2) -> inMerge l1 (c2::l2) l -> inMerge (c1::l1) (c2::l2) (c1::l)
  | inM_EQ_Some x v1 l1 v2 l2 r l: f v1 v2=Some r -> inMerge l1 l2 l -> inMerge ((x,v1)::l1) ((x,v2)::l2) ((x,r)::l)
  | inM_EQ_None x v1 l1 v2 l2 l: f v1 v2=None -> inMerge l1 l2 l -> inMerge ((x,v1)::l1) ((x,v2)::l2) l
  | inM_GT c1 l1 c2 l2 l: E.lt (fst c2) (fst c1) -> inMerge (c1::l1) l2 l -> inMerge (c1::l1) (c2::l2) (c2::l).
  Local Hint Resolve inM_nil_l inM_nir_r inM_LT inM_EQ_Some inM_EQ_None inM_GT Forall_Forall_cons.


  Lemma inMerge_app_r l2: forall l1,
     Forall (fun c => Forall (fun c' => E.lt (fst c) (fst c')) l1) l2
     -> inMerge l1 l2 (l2 ++ l1).

  Lemma inMerge_app_l l1: forall l2,
     Forall (fun c => Forall (fun c' => E.lt (fst c) (fst c')) l2) l1
     -> inMerge l1 l2 (l1 ++ l2).
  Local Hint Resolve inMerge_app_r inMerge_app_l Forall_cons.

  Lemma inMerge_GT_app l2: forall l1 r2 r3, inMerge l1 r2 r3 ->
     (Forall (fun c => Forall (fun c' => E.lt (fst c) (fst c')) l1) l2)
     -> inMerge l1 (l2 ++ r2) (l2 ++ r3).

  Lemma inMerge_LT_app l1: forall l2 r1 r3, inMerge r1 l2 r3 ->
     (Forall (fun c => Forall (fun c' => E.lt (fst c) (fst c')) l2) l1)
     -> inMerge (l1 ++ r1) l2 (l1 ++ r3).
  Local Hint Resolve inMerge_LT_app inMerge_GT_app.

  Lemma inMerge_app l1 l2 l3: inMerge l1 l2 l3 ->
    forall r1 r2 r3,
    (Forall (fun c => Forall (fun c' => E.lt (fst c) (fst c')) r2) l1)
    -> (Forall (fun c => Forall (fun c' => E.lt (fst c) (fst c')) r1) l2)
     -> inMerge r1 r2 r3
     -> inMerge (l1 ++ r1) (l2 ++ r2) (l3 ++ r3).
  Local Hint Resolve inMerge_app.

  Lemma xelements_merge m1: forall m2 path, inMerge (xelements m1 path) (xelements m2 path) (xelements (merge m1 m2) path).

  Theorem elements_merge m1 m2: inMerge (elements m1) (elements m2) (elements (merge m1 m2)).

End merge.