Require Import Bool OrderedType ZArith OrderedType OrderedTypeEx FMapInterface.

Set Implicit Arguments.
Local Open Scope positive_scope.

Fixpoint append (i j : positive) : positive :=
    match i with
      | xH => j
      | xI ii => xI (append ii j)
      | xO ii => xO (append ii j)
    end.

Lemma append_assoc_0 :
  forall (i j : positive), append i (xO j) = append (append i (xO xH)) j.

Lemma append_assoc_1 :
  forall (i j : positive), append i (xI j) = append (append i (xI xH)) j.

Lemma append_neutral_r : forall (i : positive), append i xH = i.

Lemma append_neutral_l : forall (i : positive), append xH i = i.

Module PositiveMap <: S with Module E:=PositiveOrderedTypeBits.

  Module E:=PositiveOrderedTypeBits.
  Module ME:=KeyOrderedType E.

  Definition key := positive : Type.

  #[universes(template)]
  Inductive tree (A : Type) :=
    | Leaf : tree A
    | Node : tree A -> option A -> tree A -> tree A.

  Scheme tree_ind := Induction for tree Sort Prop.

  Definition t := tree.

  Section A.
  Variable A:Type.


  Definition empty : t A := Leaf.

  Fixpoint is_empty (m : t A) : bool :=
   match m with
    | Leaf => true
    | Node l None r => (is_empty l) && (is_empty r)
    | _ => false
   end.

  Fixpoint find (i : key) (m : t A) : option A :=
    match m with
    | Leaf => None
    | Node l o r =>
        match i with
        | xH => o
        | xO ii => find ii l
        | xI ii => find ii r
        end
    end.

  Fixpoint mem (i : key) (m : t A) : bool :=
    match m with
    | Leaf => false
    | Node l o r =>
        match i with
        | xH => match o with None => false | _ => true end
        | xO ii => mem ii l
        | xI ii => mem ii r
        end
    end.

  Fixpoint add (i : key) (v : A) (m : t A) : t A :=
    match m with
    | Leaf =>
        match i with
        | xH => Node Leaf (Some v) Leaf
        | xO ii => Node (add ii v Leaf) None Leaf
        | xI ii => Node Leaf None (add ii v Leaf)
        end
    | Node l o r =>
        match i with
        | xH => Node l (Some v) r
        | xO ii => Node (add ii v l) o r
        | xI ii => Node l o (add ii v r)
        end
    end.

  Fixpoint remove (i : key) (m : t A) : t A :=
    match i with
    | xH =>
        match m with
        | Leaf => Leaf
        | Node Leaf _ Leaf => Leaf
        | Node l _ r => Node l None r
        end
    | xO ii =>
        match m with
        | Leaf => Leaf
        | Node l None Leaf =>
            match remove ii l with
            | Leaf => Leaf
            | mm => Node mm None Leaf
            end
        | Node l o r => Node (remove ii l) o r
        end
    | xI ii =>
        match m with
        | Leaf => Leaf
        | Node Leaf None r =>
            match remove ii r with
            | Leaf => Leaf
            | mm => Node Leaf None mm
            end
        | Node l o r => Node l o (remove ii r)
        end
    end.

    Fixpoint xelements (m : t A) (i : key) : list (key * A) :=
      match m with
      | Leaf => nil
      | Node l None r =>
          (xelements l (append i (xO xH))) ++ (xelements r (append i (xI xH)))
      | Node l (Some x) r =>
          (xelements l (append i (xO xH)))
          ++ ((i, x) :: xelements r (append i (xI xH)))
      end.


  Definition elements (m : t A) := xelements m xH.

  Fixpoint cardinal (m : t A) : nat :=
    match m with
      | Leaf => 0%nat
      | Node l None r => (cardinal l + cardinal r)%nat
      | Node l (Some _) r => S (cardinal l + cardinal r)
    end.

  Section CompcertSpec.

  Theorem gempty:
    forall (i: key), find i empty = None.

  Theorem gss:
    forall (i: key) (x: A) (m: t A), find i (add i x m) = Some x.

  Lemma gleaf : forall (i : key), find i (Leaf : t A) = None.

  Theorem gso:
    forall (i j: key) (x: A) (m: t A),
    i <> j -> find i (add j x m) = find i m.

  Lemma rleaf : forall (i : key), remove i Leaf = Leaf.

  Theorem grs:
    forall (i: key) (m: t A), find i (remove i m) = None.

  Theorem gro:
    forall (i j: key) (m: t A),
    i <> j -> find i (remove j m) = find i m.

  Lemma xelements_correct:
      forall (m: t A) (i j : key) (v: A),
      find i m = Some v -> List.In (append j i, v) (xelements m j).

  Theorem elements_correct:
    forall (m: t A) (i: key) (v: A),
    find i m = Some v -> List.In (i, v) (elements m).

  Fixpoint xfind (i j : key) (m : t A) : option A :=
      match i, j with
      | _, xH => find i m
      | xO ii, xO jj => xfind ii jj m
      | xI ii, xI jj => xfind ii jj m
      | _, _ => None
      end.

   Lemma xfind_left :
      forall (j i : key) (m1 m2 : t A) (o : option A) (v : A),
      xfind i (append j (xO xH)) m1 = Some v -> xfind i j (Node m1 o m2) = Some v.

    Lemma xelements_ii :
      forall (m: t A) (i j : key) (v: A),
      List.In (xI i, v) (xelements m (xI j)) -> List.In (i, v) (xelements m j).

    Lemma xelements_io :
      forall (m: t A) (i j : key) (v: A),
      ~List.In (xI i, v) (xelements m (xO j)).

    Lemma xelements_oo :
      forall (m: t A) (i j : key) (v: A),
      List.In (xO i, v) (xelements m (xO j)) -> List.In (i, v) (xelements m j).

    Lemma xelements_oi :
      forall (m: t A) (i j : key) (v: A),
      ~List.In (xO i, v) (xelements m (xI j)).

    Lemma xelements_ih :
      forall (m1 m2: t A) (o: option A) (i : key) (v: A),
      List.In (xI i, v) (xelements (Node m1 o m2) xH) -> List.In (i, v) (xelements m2 xH).

    Lemma xelements_oh :
      forall (m1 m2: t A) (o: option A) (i : key) (v: A),
      List.In (xO i, v) (xelements (Node m1 o m2) xH) -> List.In (i, v) (xelements m1 xH).

    Lemma xelements_hi :
      forall (m: t A) (i : key) (v: A),
      ~List.In (xH, v) (xelements m (xI i)).

    Lemma xelements_ho :
      forall (m: t A) (i : key) (v: A),
      ~List.In (xH, v) (xelements m (xO i)).

    Lemma find_xfind_h :
      forall (m: t A) (i: key), find i m = xfind i xH m.

    Lemma xelements_complete:
      forall (i j : key) (m: t A) (v: A),
      List.In (i, v) (xelements m j) -> xfind i j m = Some v.

  Theorem elements_complete:
    forall (m: t A) (i: key) (v: A),
    List.In (i, v) (elements m) -> find i m = Some v.

  Lemma cardinal_1 :
   forall (m: t A), cardinal m = length (elements m).

  End CompcertSpec.

  Definition MapsTo (i:key)(v:A)(m:t A) := find i m = Some v.

  Definition In (i:key)(m:t A) := exists e:A, MapsTo i e m.

  Definition Empty m := forall (a : key)(e:A) , ~ MapsTo a e m.

  Definition eq_key (p p':key*A) := E.eq (fst p) (fst p').

  Definition eq_key_elt (p p':key*A) :=
          E.eq (fst p) (fst p') /\ (snd p) = (snd p').

  Definition lt_key (p p':key*A) := E.lt (fst p) (fst p').

  Global Instance eqk_equiv : Equivalence eq_key := _.
  Global Instance eqke_equiv : Equivalence eq_key_elt := _.
  Global Instance ltk_strorder : StrictOrder lt_key := _.

  Lemma mem_find :
    forall m x, mem x m = match find x m with None => false | _ => true end.

  Lemma Empty_alt : forall m, Empty m <-> forall a, find a m = None.

  Lemma Empty_Node : forall l o r, Empty (Node l o r) <-> o=None /\ Empty l /\ Empty r.

  Section FMapSpec.

  Lemma mem_1 : forall m x, In x m -> mem x m = true.

  Lemma mem_2 : forall m x, mem x m = true -> In x m.

  Variable m m' m'' : t A.
  Variable x y z : key.
  Variable e e' : A.

  Lemma MapsTo_1 : E.eq x y -> MapsTo x e m -> MapsTo y e m.

  Lemma find_1 : MapsTo x e m -> find x m = Some e.

  Lemma find_2 : find x m = Some e -> MapsTo x e m.

  Lemma empty_1 : Empty empty.

  Lemma is_empty_1 : Empty m -> is_empty m = true.

  Lemma is_empty_2 : is_empty m = true -> Empty m.

  Lemma add_1 : E.eq x y -> MapsTo y e (add x e m).

  Lemma add_2 : ~ E.eq x y -> MapsTo y e m -> MapsTo y e (add x e' m).

  Lemma add_3 : ~ E.eq x y -> MapsTo y e (add x e' m) -> MapsTo y e m.

  Lemma remove_1 : E.eq x y -> ~ In y (remove x m).

  Lemma remove_2 : ~ E.eq x y -> MapsTo y e m -> MapsTo y e (remove x m).

  Lemma remove_3 : MapsTo y e (remove x m) -> MapsTo y e m.

  Lemma elements_1 :
     MapsTo x e m -> InA eq_key_elt (x,e) (elements m).

  Lemma elements_2 :
     InA eq_key_elt (x,e) (elements m) -> MapsTo x e m.

  Lemma xelements_bits_lt_1 : forall p p0 q m v,
     List.In (p0,v) (xelements m (append p (xO q))) -> E.bits_lt p0 p.

  Lemma xelements_bits_lt_2 : forall p p0 q m v,
     List.In (p0,v) (xelements m (append p (xI q))) -> E.bits_lt p p0.

  Lemma xelements_sort : forall p, sort lt_key (xelements m p).

  Lemma elements_3 : sort lt_key (elements m).

  Lemma elements_3w : NoDupA eq_key (elements m).

  End FMapSpec.

  Variable B : Type.

  Section Mapi.

    Variable f : key -> A -> B.

    Fixpoint xmapi (m : t A) (i : key) : t B :=
       match m with
        | Leaf => @Leaf B
        | Node l o r => Node (xmapi l (append i (xO xH)))
                             (option_map (f i) o)
                             (xmapi r (append i (xI xH)))
       end.

    Definition mapi m := xmapi m xH.

  End Mapi.

  Definition map (f : A -> B) m := mapi (fun _ => f) m.

  End A.

  Lemma xgmapi:
      forall (A B: Type) (f: key -> A -> B) (i j : key) (m: t A),
      find i (xmapi f m j) = option_map (f (append j i)) (find i m).

  Theorem gmapi:
    forall (A B: Type) (f: key -> A -> B) (i: key) (m: t A),
    find i (mapi f m) = option_map (f i) (find i m).

  Lemma mapi_1 :
    forall (elt elt':Type)(m: t elt)(x:key)(e:elt)(f:key->elt->elt'),
    MapsTo x e m ->
    exists y, E.eq y x /\ MapsTo x (f y e) (mapi f m).

  Lemma mapi_2 :
    forall (elt elt':Type)(m: t elt)(x:key)(f:key->elt->elt'),
    In x (mapi f m) -> In x m.

  Lemma map_1 : forall (elt elt':Type)(m: t elt)(x:key)(e:elt)(f:elt->elt'),
    MapsTo x e m -> MapsTo x (f e) (map f m).

  Lemma map_2 : forall (elt elt':Type)(m: t elt)(x:key)(f:elt->elt'),
    In x (map f m) -> In x m.

  Section map2.
  Variable A B C : Type.
  Variable f : option A -> option B -> option C.


  Fixpoint xmap2_l (m : t A) : t C :=
      match m with
      | Leaf => Leaf
      | Node l o r => Node (xmap2_l l) (f o None) (xmap2_l r)
      end.

  Lemma xgmap2_l : forall (i : key) (m : t A),
          f None None = None -> find i (xmap2_l m) = f (find i m) None.

  Fixpoint xmap2_r (m : t B) : t C :=
      match m with
      | Leaf => Leaf
      | Node l o r => Node (xmap2_r l) (f None o) (xmap2_r r)
      end.

  Lemma xgmap2_r : forall (i : key) (m : t B),
          f None None = None -> find i (xmap2_r m) = f None (find i m).

  Fixpoint _map2 (m1 : t A)(m2 : t B) : t C :=
    match m1 with
    | Leaf => xmap2_r m2
    | Node l1 o1 r1 =>
        match m2 with
        | Leaf => xmap2_l m1
        | Node l2 o2 r2 => Node (_map2 l1 l2) (f o1 o2) (_map2 r1 r2)
        end
    end.

    Lemma gmap2: forall (i: key)(m1:t A)(m2: t B),
      f None None = None ->
      find i (_map2 m1 m2) = f (find i m1) (find i m2).

  End map2.

  Definition map2 (elt elt' elt'':Type)(f:option elt->option elt'->option elt'') :=
   _map2 (fun o1 o2 => match o1,o2 with None,None => None | _, _ => f o1 o2 end).

  Lemma map2_1 : forall (elt elt' elt'':Type)(m: t elt)(m': t elt')
    (x:key)(f:option elt->option elt'->option elt''),
    In x m \/ In x m' ->
    find x (map2 f m m') = f (find x m) (find x m').

  Lemma map2_2 : forall (elt elt' elt'':Type)(m: t elt)(m': t elt')
    (x:key)(f:option elt->option elt'->option elt''),
    In x (map2 f m m') -> In x m \/ In x m'.

  Section Fold.

    Variables A B : Type.
    Variable f : key -> A -> B -> B.

    Fixpoint xfoldi (m : t A) (v : B) (i : key) :=
      match m with
        | Leaf _ => v
        | Node l (Some x) r =>
          xfoldi r (f i x (xfoldi l v (append i 2))) (append i 3)
        | Node l None r =>
          xfoldi r (xfoldi l v (append i 2)) (append i 3)
      end.

    Lemma xfoldi_1 :
      forall m v i,
      xfoldi m v i = fold_left (fun a p => f (fst p) (snd p) a) (xelements m i) v.

    Definition fold m i := xfoldi m i 1.

  End Fold.

  Lemma fold_1 :
    forall (A:Type)(m:t A)(B:Type)(i : B) (f : key -> A -> B -> B),
    fold f m i = fold_left (fun a p => f (fst p) (snd p) a) (elements m) i.

  Fixpoint equal (A:Type)(cmp : A -> A -> bool)(m1 m2 : t A) : bool :=
    match m1, m2 with
      | Leaf _, _ => is_empty m2
      | _, Leaf _ => is_empty m1
      | 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 cmp l1 l2 && equal cmp r1 r2
     end.

  Definition Equal (A:Type)(m m':t A) :=
    forall y, find y m = find y m'.
  Definition Equiv (A:Type)(eq_elt:A->A->Prop) m m' :=
    (forall k, In k m <-> In k m') /\
    (forall k e e', MapsTo k e m -> MapsTo k e' m' -> eq_elt e e').
  Definition Equivb (A:Type)(cmp: A->A->bool) := Equiv (Cmp cmp).

  Lemma equal_1 : forall (A:Type)(m m':t A)(cmp:A->A->bool),
    Equivb cmp m m' -> equal cmp m m' = true.

  Lemma equal_2 : forall (A:Type)(m m':t A)(cmp:A->A->bool),
    equal cmp m m' = true -> Equivb cmp m m'.

End PositiveMap.

Module PositiveMapAdditionalFacts.
  Import PositiveMap.

  Theorem gsspec:
    forall (A:Type)(i j: key) (x: A) (m: t A),
    find i (add j x m) = if E.eq_dec i j then Some x else find i m.

  Theorem gsident:
    forall (A:Type)(i: key) (m: t A) (v: A),
    find i m = Some v -> add i v m = m.

  Lemma xmap2_lr :
      forall (A B : Type)(f g: option A -> option A -> option B)(m : t A),
      (forall (i j : option A), f i j = g j i) ->
      xmap2_l f m = xmap2_r g m.

  Theorem map2_commut:
    forall (A B: Type) (f g: option A -> option A -> option B),
    (forall (i j: option A), f i j = g j i) ->
    forall (m1 m2: t A),
    _map2 f m1 m2 = _map2 g m2 m1.

End PositiveMapAdditionalFacts.