include "basics/types.ma".
include "utilities/binary/positive.ma".
include "ASM/Util.ma". (* bool_to_Prop *)

inductive positive_map (A:Type[0]) : Type[0] ≝
| pm_leaf : positive_map A
| pm_node : option A → positive_map A → positive_map A → positive_map A.

let rec lookup_opt (A:Type[0]) (b:Pos) (t:positive_map A) on t : option A ≝
match t with
[ pm_leaf ⇒ None ?
| pm_node a l r ⇒
    match b with
    [ one ⇒ a
    | p0 tl ⇒ lookup_opt A tl l
    | p1 tl ⇒ lookup_opt A tl r
    ]
].

definition lookup: ∀A:Type[0].Pos → positive_map A → A → A ≝  
 λA.λb.λt.λx.
 match lookup_opt A b t with
 [ None ⇒ x
 | Some r ⇒ r
 ].

let rec pm_set (A:Type[0]) (b:Pos) (a:option A) (t:positive_map A) on b : positive_map A ≝
match b with
[ one ⇒
    match t with
    [ pm_leaf ⇒ pm_node A a (pm_leaf A) (pm_leaf A)
    | pm_node _ l r ⇒ pm_node A a l r
    ]
| p0 tl ⇒
    match t with
    [ pm_leaf ⇒ pm_node A (None ?) (pm_set A tl a (pm_leaf A)) (pm_leaf A)
    | pm_node x l r ⇒ pm_node A x (pm_set A tl a l) r
    ]
| p1 tl ⇒
    match t with
    [ pm_leaf ⇒ pm_node A (None ?) (pm_leaf A) (pm_set A tl a (pm_leaf A))
    | pm_node x l r ⇒ pm_node A x l (pm_set A tl a r)
    ]
].

definition insert : ∀A:Type[0]. Pos → A → positive_map A → positive_map A ≝
λA,p,a. pm_set A p (Some ? a).

let rec update (A:Type[0]) (b:Pos) (a:A) (t:positive_map A) on b : option (positive_map A) ≝
match b with
[ one ⇒
    match t with
    [ pm_leaf ⇒ None ?
    | pm_node x l r ⇒ option_map ?? (λ_. pm_node A (Some ? a) l r) x
    ]
| p0 tl ⇒
    match t with
    [ pm_leaf ⇒ None ?
    | pm_node x l r ⇒ option_map ?? (λl. pm_node A x l r) (update A tl a l)
    ]
| p1 tl ⇒
    match t with
    [ pm_leaf ⇒ None ?
    | pm_node x l r ⇒ option_map ?? (λr. pm_node A x l r) (update A tl a r)
    ]
].

lemma lookup_opt_pm_set_hit :
 ∀A:Type[0].∀v:option A.∀b:Pos.∀t:positive_map A.
  lookup_opt … b (pm_set … b v t) = v.
#A #v #b elim b
[ * //
| #tl #IH * 
  [ whd in ⊢ (??%%); @IH
  | #x #l #r @IH
  ]
| #tl #IH * 
  [ whd in ⊢ (??%%); @IH
  | #x #l #r @IH
  ]
] qed.

lemma lookup_opt_insert_hit : 
 ∀A:Type[0].∀v:A.∀b:Pos.∀t:positive_map A.
  lookup_opt … b (insert … b v t) = Some A v.
#A #v #b elim b
[ * //
| #tl #IH * 
  [ whd in ⊢ (??%%); @IH
  | #x #l #r @IH
  ]
| #tl #IH * 
  [ whd in ⊢ (??%%); @IH
  | #x #l #r @IH
  ]
] qed.

lemma lookup_insert_hit:
  ∀a: Type[0].
  ∀v: a.
  ∀b: Pos.
  ∀t: positive_map a.
  ∀d: a.
    lookup … b (insert … b v t) d = v.
#A #v #b #t #d whd in ⊢ (??%?); >lookup_opt_insert_hit %
qed.

lemma lookup_opt_pm_set_miss:
 ∀A:Type[0].∀v:option A.∀b,c:Pos.∀t:positive_map A.
  b ≠ c → lookup_opt … b (pm_set … c v t) = lookup_opt … b t.
#A #v #b elim b
[ * [ #t * #H elim (H (refl …))
    | *: #c' #t #NE cases t //
    ]
| #b' #IH *
  [ * [ #NE @refl | #x #l #r #NE @refl ]
  | #c' * [ #NE whd in ⊢ (??%%); @IH /2/
          | #x #l #r #NE whd in ⊢ (??%%); @IH /2/
          ]
  | #c' * //
  ]
| #b' #IH *
  [ * [ #NE @refl | #x #l #r #NE @refl ]
  | #c' * //
  | #c' * [ #NE whd in ⊢ (??%%); @IH /2/
          | #x #l #r #NE whd in ⊢ (??%%); @IH /2/
          ]
  ]
] qed.

lemma lookup_opt_insert_miss:
 ∀A:Type[0].∀v:A.∀b,c:Pos.∀t:positive_map A.
  b ≠ c → lookup_opt … b (insert … c v t) = lookup_opt … b t.
#A #v #b elim b
[ * [ #t * #H elim (H (refl …))
    | *: #c' #t #NE cases t //
    ]
| #b' #IH *
  [ * [ #NE @refl | #x #l #r #NE @refl ]
  | #c' * [ #NE whd in ⊢ (??%%); @IH /2/
          | #x #l #r #NE whd in ⊢ (??%%); @IH /2/
          ]
  | #c' * //
  ]
| #b' #IH *
  [ * [ #NE @refl | #x #l #r #NE @refl ]
  | #c' * //
  | #c' * [ #NE whd in ⊢ (??%%); @IH /2/
          | #x #l #r #NE whd in ⊢ (??%%); @IH /2/
          ]
  ]
] qed.

let rec fold (A, B: Type[0]) (f: Pos → A → B → B)
 (t: positive_map A) (b: B) on t: B ≝
match t with
[ pm_leaf ⇒ b
| pm_node a l r ⇒
    let b ≝ match a with [ None ⇒ b | Some a ⇒ f one a b ] in
    let b ≝ fold A B (λx.f (p0 x)) l b in
            fold A B (λx.f (p1 x)) r b
].    

definition domain_of_pm : ∀A. positive_map A → positive_map unit ≝
λA,t. fold A (positive_map unit) (λp,a,b. insert unit p it b) t (pm_leaf unit).

(* Build some results about fold using domain_of_pm. *)

lemma pm_fold_ignore_adding_junk : ∀A. ∀m. ∀f,g:Pos → Pos. ∀s.
  (∀p,q. f p ≠ g q) →
  ∀p. lookup_opt unit (f p) (fold A ? (λx,a,b. insert unit (g x) it b) m s) = lookup_opt unit (f p) s.
#A #m elim m
[ normalize //
| #oa #l #r #IHl #IHr
  #f #g #s #not_g #p
  normalize
  >IHr
  [ >IHl
    [ cases oa
      [ normalize %
      | #a normalize >(lookup_opt_insert_miss ?? (f p)) [ % | @not_g ]
      ]
    | #x #y % #E cases (not_g x (p0 y)) #H @H @E
    ]
  | #x #y % #E cases (not_g x (p1 y)) #H @H @E
  ]
] qed.

lemma pm_fold_ind_gen : ∀A,B,m,f,b.
  ∀pre:Pos→Pos. (∀x,y.pre x = pre y → x = y) → ∀ps. (∀p. lookup_opt ? (pre p) ps = None ?) →
  ∀P:B → positive_map unit → Prop.
  P b ps →
  (∀p,ps,a,b. lookup_opt unit (pre p) ps = None unit → lookup_opt A p m = Some A a → P b ps → P (f (pre p) a b) (insert unit (pre p) it ps)) →
  P (fold A B (λx. f (pre x)) m b) (fold ?? (λp,a,b. insert unit (pre p) it b) m ps).
#A #B #m elim m
[ //
| #oa #l #r #IHl #IHr
  #f #b #pre #PRE #ps #PS #P #emp #step whd in ⊢ (?%%);
  @IHr
  [ #x #y #E lapply (PRE … E) #E' destruct //
  | #p change with ((λx. pre (p1 x)) p) in match (pre (p1 p)); >pm_fold_ignore_adding_junk
    [ cases oa [ @PS | #a >(lookup_opt_insert_miss ?? (pre (p1 p))) [ @PS | % #E lapply (PRE … E) #E' destruct ] ]
    | #p #q % #E lapply (PRE … E) #E' destruct
    ]
  | @IHl
    [ #x #y #E lapply (PRE … E) #E' destruct //
    | #p cases oa [ @PS | #a >(lookup_opt_insert_miss ?? (pre (p0 p))) [ @PS | % #E lapply (PRE … E) #E' destruct ] ]
    | cases oa in step ⊢ %; [ #_ @emp | #a #step normalize @step [ @PS | % | @emp ] ]
    | #p #ps' #a #b' @step
    ]
  | #p #ps' #a #b' @step
  ]
] qed.

lemma pm_fold_ext : ∀A,B,m,f,b.
  fold A B f m b = fold A B (λx:Pos. f x) m b.
#A #B #m elim m /4/
qed.

(* Main result about folds.  (Rather than producing a result about just the
   domain of m we could use the whole mapping, but we'd need a function to
   canonicalise maps because their representation can contain some junk.) *)

lemma pm_fold_ind : ∀A,B,f,m,b. ∀P:B → positive_map unit → Prop.
  P b (pm_leaf unit) →
  (∀p,ps,a,b. lookup_opt unit p ps = None unit → lookup_opt A p m = Some A a → P b ps → P (f p a b) (insert unit p it ps)) →
  P (fold A B f m b) (domain_of_pm A m).
#A #B #f #m #b #P #emp #step
>pm_fold_ext
whd in ⊢ (??%);
@(pm_fold_ind_gen A B m f b (λx.x) ? (pm_leaf unit) ???)
[ //
| //
| @emp
| @step
] qed.

lemma pm_fold_eq : ∀A,B,b,f.
  (∀p,k:Pos. ∀a,s1,s2. lookup_opt B p s1 = lookup_opt B p s2 → lookup_opt B p (f k a s1) = lookup_opt B p (f k a s2)) →
  ∀p,s1,s2.
  lookup_opt B p s1 = lookup_opt B p s2 →
  lookup_opt B p (fold A (positive_map B) f b s1) = lookup_opt B p (fold A (positive_map B) f b s2).
#A #B #b elim b
[ #f #H #p #s1 #s2 #H @H
| #oa #l #r #IHl #IHr #f #H #p #s1 #s2 #H'
  whd in match (fold ???? s1);
  whd in match (fold ???? s2);
  @IHr [ #q #k #a #s1' #s2' #H'' @H @H'' ]
  @IHl [ #q #k #a #s1' #s2' #H'' @H @H'' ]
  cases oa [ @H' | #a @H @H' ]
] qed.


lemma pm_fold_ignore_prefix : ∀A. ∀m. ∀pre:Pos → Pos. (∀x,y. pre x = pre y → x = y) →
  ∀p. ∀pre'.
      lookup_opt unit (pre p) (fold A ? (λx,a,b. insert unit (pre (pre' x)) it b) m (pm_leaf ?)) =
      lookup_opt unit p (fold A ? (λx,a,b. insert unit (pre' x) it b) m (pm_leaf ?)).
#A #m #pre #PRE
cut (∀p. lookup_opt unit (pre p) (pm_leaf unit) = lookup_opt unit p (pm_leaf unit)) [ // ]
generalize in match (pm_leaf unit) in ⊢ ((? → ??%?) → ? → ? → ??%?);
generalize in match (pm_leaf unit);
elim m
[ #s1 #s2 #S #p #pre' @S 
| #oa #l #r #IHl #IHr
  #s1 #s2 #S #p #pre'
  whd in ⊢ (??(???%)(???%));
  @(IHr ???? (λx. pre' (p1 x)))
  #p' @IHl
  #p'' cases oa
  [ @S
  | #a cases (decidable_eq_pos p'' (pre' one))
    [ #E destruct
      >(lookup_opt_insert_hit ?? (pre (pre' one)))
      >(lookup_opt_insert_hit ?? (pre' one))
      %
    | #NE >(lookup_opt_insert_miss ??? (pre (pre' one)))
      [ >(lookup_opt_insert_miss ??? (pre' one)) //
      | % #E @(absurd … (PRE … E) NE)
      ]
    ]
  ]
] qed.

(* Link the domain of a map (used in the result about fold) to the original
   map. *)

lemma domain_of_pm_present : ∀A,m,p.
  lookup_opt A p m ≠ None ? ↔ lookup_opt unit p (domain_of_pm A m) ≠ None ?.

#A #m whd in match (domain_of_pm ??);
elim m
[ #p normalize /3/
| #oa #l #r #IHl #IHr
  whd in match (fold ???? (pm_leaf unit));
  *
  [ change with ((λx:Pos. one) one) in ⊢ (??(?(??(??%?)?)));
    >pm_fold_ignore_adding_junk [2: #p #q % #E destruct ]
    change with ((λx:Pos. one) one) in ⊢ (??(?(??(??%?)?)));
    >pm_fold_ignore_adding_junk [2: #p #q % #E destruct ]
    cases oa
    [ normalize /3/
    | #a change with (insert unit one ??) in ⊢ (??(?(??(???%)?)));
      >(lookup_opt_insert_hit ?? one) normalize
      % #_ % #E destruct
    ]
  | #p >(pm_fold_eq ??????? (pm_leaf unit))
    [ >pm_fold_ignore_prefix [2: #x #y #E destruct % ]
      @IHr
    | >pm_fold_ignore_adding_junk [2: #x #y % #E destruct ]
      cases oa
      [ %
      | #a >(lookup_opt_insert_miss ?? (p1 p)) [2: % #E destruct ]
        %
      ]
    | #x #y #a #s1 #s2 #S cases (decidable_eq_pos x (p1 y))
      [ #E destruct
        >(lookup_opt_insert_hit ?? (p1 y))
        >(lookup_opt_insert_hit ?? (p1 y))
        %
      | #NE >(lookup_opt_insert_miss ?? x … NE)
        >(lookup_opt_insert_miss ?? x … NE)
        @S
      ]
    ]  
  | #p
    >pm_fold_ignore_adding_junk [2: #x #y % #E destruct ]
    >(pm_fold_eq ??????? (pm_leaf unit))
    [ >pm_fold_ignore_prefix [2: #x #y #E destruct % ]
      @IHl
    | cases oa
      [ %
      | #a >(lookup_opt_insert_miss ?? (p0 p)) [2: % #E destruct ]
        %
      ]
    | #x #y #a #s1 #s2 #S cases (decidable_eq_pos x (p0 y))
      [ #E destruct
        >(lookup_opt_insert_hit ?? (p0 y))
        >(lookup_opt_insert_hit ?? (p0 y))
        %
      | #NE >(lookup_opt_insert_miss ?? x … NE)
        >(lookup_opt_insert_miss ?? x … NE)
        @S
      ]
    ]
  ]
] qed.


lemma update_fail : ∀A,b,a,t.
  update A b a t = None ? →
  lookup_opt A b t = None ?.
#A #b #a elim b
[ * [ // | * // #x #l #r normalize #E destruct ]
| #b' #IH * [ // | #x #l #r #U @IH whd in U:(??%?);
                   cases (update A b' a r) in U ⊢ %;
                   [ // | normalize #t #E destruct ]
            ]
| #b' #IH * [ // | #x #l #r #U @IH whd in U:(??%?);
                   cases (update A b' a l) in U ⊢ %;
                   [ // | normalize #t #E destruct ]
            ]
] qed.

lemma update_lookup_opt_same : ∀A,b,a,t,t'.
  update A b a t = Some ? t' →
  lookup_opt A b t' = Some ? a.
#A #b #a elim b
[ * [ #t' normalize #E destruct
    | * [ 2: #x ] #l #r #t' whd in ⊢ (??%? → ??%?);
      #E destruct //
    ]
| #b' #IH *
  [ #t' normalize #E destruct
  | #x #l #r #t' whd in ⊢ (??%? → ??%?);
    lapply (IH r)
    cases (update A b' a r)
    [ #_ normalize #E destruct
    | #r' #H normalize #E destruct @H @refl
    ]
  ]
| #b' #IH *
  [ #t' normalize #E destruct
  | #x #l #r #t' whd in ⊢ (??%? → ??%?);
    lapply (IH l)
    cases (update A b' a l)
    [ #_ normalize #E destruct
    | #r' #H normalize #E destruct @H @refl
    ]
  ]
] qed.

lemma update_lookup_opt_other : ∀A,b,a,t,t'.
  update A b a t = Some ? t' →
  ∀b'. b ≠ b' →
  lookup_opt A b' t = lookup_opt A b' t'.
#A #b #a elim b
[ * [ #t' normalize #E destruct
    | * [2: #x] #l #r #t' normalize #E destruct
      * [ * #H elim (H (refl …)) ]
      #tl #NE normalize //
    ]
| #b' #IH *
  [ #t' normalize #E destruct
  | #x #l #r #t' normalize
    lapply (IH r)
    cases (update A b' a r)
    [ #_ normalize #E destruct
    | #t'' #H normalize #E destruct * // #b'' #NE @H /2/
    ]
  ] 
| #b' #IH *
  [ #t' normalize #E destruct
  | #x #l #r #t' normalize
    lapply (IH l)
    cases (update A b' a l)
    [ #_ normalize #E destruct
    | #t'' #H normalize #E destruct * // #b'' #NE @H /2/
    ]
  ]
] qed.

include "utilities/option.ma".

let rec map_opt A B (f : A → option B) (b : positive_map A) on b : positive_map B ≝
match b with
[ pm_leaf ⇒ pm_leaf ?
| pm_node a l r ⇒ pm_node ? (!x ← a ; f x)
    (map_opt ?? f l)
    (map_opt ?? f r)
].

definition map ≝ λA,B,f.map_opt A B (λx. Some ? (f x)).

lemma lookup_opt_map : ∀A,B,f,b,p.
  lookup_opt … p (map_opt A B f b) = ! x ← lookup_opt … p b ; f x.
#A#B#f#b elim b [//]
#a #l #r #Hil #Hir * [//]
#p
whd in ⊢ (??(???%)(????%?));
whd in ⊢ (??%?); //
qed.

lemma map_opt_id_eq_ext : ∀A,m,f.(∀x.f x = Some ? x) → map_opt A A f m = m.
#A #m #f #Hf elim m [//]
* [2:#x] #l #r #Hil #Hir normalize [>Hf] //
qed.

lemma map_opt_id : ∀A,m.map_opt A A (λx. Some ? x) m = m.
#A #m @map_opt_id_eq_ext //
qed.

let rec merge A B C (choice : option A → option B → option C)
  (a : positive_map A) (b : positive_map B) on a : positive_map C ≝
match a with
[ pm_leaf ⇒ map_opt ?? (λx.choice (None ?) (Some ? x)) b
| pm_node o l r ⇒
  match b with
  [ pm_leaf ⇒ map_opt ?? (λx.choice (Some ? x) (None ?)) a
  | pm_node o' l' r' ⇒
    pm_node ? (choice o o') 
      (merge … choice l l')
      (merge … choice r r')
  ]
].

lemma lookup_opt_merge : ∀A,B,C,choice,a,b,p.
  choice (None ?) (None ?) = None ? →
  lookup_opt … p (merge A B C choice a b) =
    choice (lookup_opt … p a) (lookup_opt … p b).
 #A#B#C#choice#a elim a
[ normalize #b
| #o#l#r#Hil#Hir * [2: #o'#l'#r' * normalize /2 by /]
]
#p#choice_good >lookup_opt_map
elim (lookup_opt ???)
normalize //
qed.

let rec domain_size (A : Type[0]) (b : positive_map A) on b : ℕ ≝
match b with
[ pm_leaf ⇒ 0
| pm_node a l r ⇒
  (match a with
   [ Some _ ⇒ 1
   | None ⇒ 0
   ]) + domain_size A l + domain_size A r
].

(* just to load notation for 'norm (list.ma's notation overrides core's 'card)*)
include "basics/lists/list.ma".
interpretation "positive map size" 'norm p = (domain_size ? p).

lemma map_opt_size : ∀A,B,f,b.|map_opt A B f b| ≤ |b|.
#A#B#f#b elim b [//]
*[2:#x]#l#r#Hil#Hir normalize
[elim (f ?) normalize [@(transitive_le ? ? ? ? (le_n_Sn ?)) | #y @le_S_S ] ]
@le_plus assumption
qed.

lemma map_size : ∀A,B,f,b.|map A B f b| = |b|.
#A#B#f#b elim b [//]
*[2:#x]#l#r#Hil#Hir normalize
>Hil >Hir @refl
qed.

lemma lookup_opt_O_lt_size : ∀A,m,p. lookup_opt A p m ≠ None ? → 0 < |m|.
#A#m elim m
[#p normalize #F elim (absurd ? (refl ??) F)
|* [2: #x]  #l #r #Hil #Hir * normalize
  [1,2,3: //
  |4:  #F elim (absurd ? (refl ??) F)]
  #p #H [@(transitive_le … (Hir ? H)) | @(transitive_le … (Hil ? H))] //
qed.

lemma insert_size : ∀A,p,a,m. |insert A p a m| =
  (match lookup_opt … p m with [Some _ ⇒ 0 | None ⇒ 1]) + |m|.
#A#p elim p
[ #a * [2: * [2: #x] #l #r] normalize //
|*: #p #Hi #a * [2,4: * [2,4: #x] #l #r] normalize >Hi //
] qed.

(* Testing all of the entries with a boolean predicate, where the predicate
   is given a proof that the entry is in the map. *)

let rec pm_all_aux (A:Type[0]) (m,t:positive_map A) (pre:Pos → Pos) on t : (∀p. lookup_opt A (pre p) m = lookup_opt A p t) → (∀p,a.lookup_opt A p m = Some A a → bool) → bool ≝
match t return λt. (∀p. lookup_opt A (pre p) m = lookup_opt A p t) → (∀p,a.lookup_opt A p m = Some A a → bool) → bool with
[ pm_leaf ⇒ λ_.λ_. true
| pm_node a l r ⇒ λH,f.
    pm_all_aux A m l (λx. pre (p0 x)) ? f ∧
    ((match a return λa. (∀p. ? = lookup_opt A p (pm_node ? a ??)) → ? with [ None ⇒ λ_. true | Some a' ⇒ λH. f (pre one) a' ? ]) H) ∧
    pm_all_aux A m r (λx. pre (p1 x)) ? f
].
[ >H %
| #p >H %
| #p >H %
] qed.

definition pm_all : ∀A. ∀m:positive_map A. (∀p,a. lookup_opt A p m = Some A a → bool) → bool ≝
λA,m,f. pm_all_aux A m m (λx.x) (λp. refl ??) f.

(* Proof irrelevance doesn't seem to apply to arbitrary variables, but we can
   use the function graph to show that it's fine. *)
   
inductive pm_all_pred_graph : ∀A,m,p,a,L. ∀f:∀p:Pos.∀a:A.lookup_opt A p m=Some A a→bool. bool → Prop ≝
| pmallpg : ∀A,m,p,a,L. ∀f:∀p:Pos.∀a:A.lookup_opt A p m=Some A a→bool. pm_all_pred_graph A m p a L f (f p a L).

lemma pm_all_pred_irr : ∀A,m,p,a,L,L'.
 ∀f:∀p:Pos.∀a:A.lookup_opt A p m=Some A a→bool.
 f p a L = f p a L'.
#A #m #p #a #L #L' #f
cut (pm_all_pred_graph A m p a L f (f p a L)) [ % ]
cases (f p a L) #H cut (pm_all_pred_graph A m p a L' f ?) [2,5: @H | 1,4: skip]
* //
qed.

lemma pm_all_aux_ok_aux : ∀A,m,t,pre,H,f,a,L.
  pm_all_aux A m t pre H f →
  lookup_opt A one t = Some A a →
  f (pre one) a L.
#A #m *
[ #pre #H #f #a #L #H #L' normalize in L'; destruct
| * [ #l #r #pre #H #f #a #L #H' #L' normalize in L'; destruct
    | #a' #l #r #pre #H #f #a #L #AUX #L' normalize in L'; destruct
      cases (andb_Prop_true … AUX) #AUX' cases (andb_Prop_true … AUX')
      #_ #HERE #_ whd in HERE:(?%); @eq_true_to_b <HERE @pm_all_pred_irr
    ]
] qed.

lemma pm_all_aux_ok : ∀A,m,t,pre,H,f.
  bool_to_Prop (pm_all_aux A m t pre H f) ↔ (∀p,a,H. f (pre p) a H).
#A #m #t #pre #H #f %
[ #H' #p generalize in match pre in H H' ⊢ %; -pre generalize in match t; elim p
  [ #t #pre #H #H' #a #L lapply (refl ? (Some A a)) <L in ⊢ (??%? → ?); >H @pm_all_aux_ok_aux //
  | #q #IH * [ #pre #H #H' #a #L @⊥ >H in L; normalize #E destruct
             | #x #l #r #pre #H #H' #a #L @(IH r)
               [ #x >H % | cases (andb_Prop_true … H') #_ // ]
             ]
  | #q #IH * [ #pre #H #H' #a #L @⊥ >H in L; normalize #E destruct
             | #x #l #r #pre #H #H' #a #L @(IH l)
               [ #x >H % | cases (andb_Prop_true … H') #H'' cases (andb_Prop_true … H'') // ]
             ]
  ]
| #H' generalize in match pre in H H' ⊢ %; -pre elim t
  [ //
  | * [2:#x] #l #r #IHl #IHr #pre #H #H' @andb_Prop
    [ 1,3: @andb_Prop
      [ 1,3: @IHl //
      | @(H' one x)
      | %
      ]
    | 2,4: @IHr //
    ]
  ]
] qed.

lemma pm_all_ok : ∀A,m,f.
  bool_to_Prop (pm_all A m f) ↔ (∀p,a,H. f p a H).
#A #m #f @pm_all_aux_ok
qed.


(* Attempt to choose an entry in the map, and if successful return the entry
   and the map without it. *)

let rec pm_choose (A:Type[0]) (t:positive_map A) on t : option (Pos × A × (positive_map A)) ≝
match t with
[ pm_leaf ⇒ None ?
| pm_node a l r ⇒
  match pm_choose A l with
  [ None ⇒
    match pm_choose A r with
    [ None ⇒
      match a with
      [ None ⇒ None ?
      | Some a ⇒ Some ? 〈〈one, a〉, pm_leaf A〉
      ]
    | Some x ⇒ Some ? 〈〈p1 (\fst (\fst x)), \snd (\fst x)〉, pm_node A a l (\snd x)〉
    ]
  | Some x ⇒ Some ? 〈〈p0 (\fst (\fst x)), \snd (\fst x)〉, pm_node A a (\snd x) r〉
  ]
].

lemma pm_choose_empty : ∀A,t.
  pm_choose A t = None ? ↔ ∀p. lookup_opt A p t = None A.
#A #t elim t
[ % //
| *
  [ #l #r * #IHl #IHl' * #IHr #IHr' %
    [ #C *
      [ %
      | #p whd in C:(??%?) ⊢ (??%?);
        @IHr cases (pm_choose A r) in C ⊢ %;
        [ //
        | #x normalize cases (pm_choose A l) [2: #y] normalize #E destruct
        ] 
      | #p whd in C:(??%?) ⊢ (??%?);
        @IHl cases (pm_choose A l) in C ⊢ %;
        [ //
        | #x normalize #E destruct
        ]
      ]
    | #L whd in ⊢ (??%?); >IHl' [ normalize >IHr' [ % | #p @(L (p1 p)) ] | #p @(L (p0 p)) ]
    ]
  | #a #l #r #IHl #IHr %
    [ normalize cases (pm_choose A l) [2: normalize #x #E destruct ]
      normalize cases (pm_choose A r) [2: #x] normalize #E destruct
    | #L lapply (L one) normalize #E destruct
    ]
  ]
] qed.

lemma pm_choose_empty_card : ∀A,t.
  pm_choose A t = None ? ↔ |t| = 0.
#A #t elim t
[ /2/
| * [ #l #r * #IHl #IHl' * #IHr #IHr' %
      [ normalize lapply IHl cases (pm_choose A l) [2: #x #_ #E normalize in E; destruct ]
        normalize lapply IHr cases (pm_choose A r) [2: #x #_ #_ #E normalize in E; destruct ]
        #H1 #H2 #_ >(H1 (refl ??)) >(H2 (refl ??)) %
      | normalize #CARD >IHl' [ >IHr' ] /2 by refl, le_n_O_to_eq/
      ]
    | #a #l #r #_ #_ % normalize
      [ cases (pm_choose A l) [ cases (pm_choose A r) [ | #x ] | #x ]
        #E normalize in E; destruct
      | #E destruct
      ]
    ]
] qed.

lemma pm_choose_some : ∀A,t,p,a,t'.
  pm_choose A t = Some ? 〈〈p,a〉,t'〉 →
  lookup_opt A p t = Some A a ∧
  lookup_opt A p t' = None A ∧
  (∀q. p = q ∨ lookup_opt A q t = lookup_opt A q t') ∧
  |t| = S (|t'|).
#A #t elim t
[ #p #a #t' normalize #E destruct
| #ao #l #r #IHl #IHr *
  [ #a #t' normalize
    lapply (pm_choose_empty_card A l)
    lapply (pm_choose_empty A l)
    cases (pm_choose A l) normalize [2: #x #_ #_ #E destruct ] * #EMPl #EMPl' * #CARDl #CARDl'
    lapply (pm_choose_empty_card A r)
    lapply (pm_choose_empty A r)
    cases (pm_choose A r) normalize [2: #x #_ #_ #E destruct ] * #EMPr #EMPr' * #CARDr #CARDr'
    cases ao normalize [2:#x] #E destruct
    % [ % /2/ * /2/ #q %2 whd in ⊢ (??%%); /2/ | >CARDl >CARDr // ]
  | #p #a #t' normalize
    lapply (pm_choose_empty_card A l)
    lapply (pm_choose_empty A l)
    cases (pm_choose A l) normalize [2: #x #_ #_ #E destruct ] * #EMPl #EMPl' * #CARDl #CARDl'
    lapply (IHr p) 
    cases (pm_choose A r) normalize [ #_ cases ao [2:#a'] normalize #E destruct ]
    * * #p' #a' #t'' #IH #E destruct
    cases (IH ?? (refl ??)) * * #L1 #L2 #L3 #CARD
    % [ % [ % [ // | @L2 ]| * /2/ #q cases (L3 q) /2/ #L4 %2 @L4 ] | cases ao normalize >CARD // ]
  | #p #a #t' normalize
    lapply (IHl p) 
    cases (pm_choose A l) normalize
    [ #_ cases (pm_choose A r) normalize [ cases ao [2:#a'] normalize #E destruct
      | #x #E destruct ]
    |  * * #p' #a' #t'' #IH #E destruct
      cases (IH ?? (refl ??)) * * #L1 #L2 #L3 #CARD
      % [ % [ % [ // | @L2 ]| * /2/ #q cases (L3 q) /2/ #L4 %2 @L4 ] | cases ao normalize >CARD // ]
    ]
  ]
] qed.

(* Try to remove an element, return updated map if successful *)

let rec pm_try_remove (A:Type[0]) (b:Pos) (t:positive_map A) on b : option (A × (positive_map A)) ≝
match b with
[ one ⇒
    match t with
    [ pm_leaf ⇒ None ?
    | pm_node x l r ⇒ option_map ?? (λx. 〈x, pm_node A (None ?) l r〉) x
    ]
| p0 tl ⇒
    match t with
    [ pm_leaf ⇒ None ?
    | pm_node x l r ⇒ option_map ?? (λxl. 〈\fst xl, pm_node A x (\snd xl) r〉) (pm_try_remove A tl l)
    ]
| p1 tl ⇒
    match t with
    [ pm_leaf ⇒ None ?
    | pm_node x l r ⇒ option_map ?? (λxr. 〈\fst xr, pm_node A x l (\snd xr)〉) (pm_try_remove A tl r)
    ]
].

lemma pm_try_remove_none : ∀A,b,t.
  pm_try_remove A b t = None ? ↔ lookup_opt A b t = None ?.
#A #b elim b
[ * [ /2/ | * [ /2/ | #a #l #r % normalize #E destruct ] ]
| #p #IH * [ /2/ | #x #l #r % normalize cases (IH r) cases (pm_try_remove A p r)
  [ normalize #H #X @H | * #a #t' #_ normalize #_ #X destruct
  | normalize // | * #a #t' normalize #A #B #C lapply (B C) #E destruct
  ] ]
| #p #IH * [ /2/ | #x #l #r % normalize cases (IH l) cases (pm_try_remove A p l)
  [ normalize #H #X @H | * #a #t' #_ normalize #_ #X destruct
  | normalize // | * #a #t' normalize #A #B #C lapply (B C) #E destruct
  ] ]
] qed.

lemma pm_try_remove_some : ∀A,p,t,a,t'.
  pm_try_remove A p t = Some ? 〈a,t'〉 →
  lookup_opt A p t = Some A a ∧
  lookup_opt A p t' = None A ∧
  (∀q. p = q ∨ lookup_opt A q t = lookup_opt A q t') ∧
  |t| = S (|t'|).
#A #p elim p
[ * [ #a #t' #E normalize in E; destruct
    | * [ #l #r #a #t' #E normalize in E; destruct
        | #a' #l #r #a #t' #E normalize in E; destruct % [ % [ % // | * /2/ ]| // ]
        ]
    ]
| #p' #IH *
  [ #a #t' #E normalize in E; destruct
  | #x #l #r #a #t' whd in ⊢ (??%? → ?);
    lapply (IH r) cases (pm_try_remove A p' r) [ #_ #E normalize in E; destruct ]
    * #a' #r' #H #E whd in E:(??%?); destruct
    cases (H a r' (refl ??)) * * #L1 #L2 #L3 #CARD
    @conj try @conj try @conj
    [ @L1 | @L2 | * /2/ #q cases (L3 q) [ /2/ | #E %2 @E ] | cases x normalize >CARD // ]
  ]
| #p' #IH *
  [ #a #t' #E normalize in E; destruct
  | #x #l #r #a #t' whd in ⊢ (??%? → ?);
    lapply (IH l) cases (pm_try_remove A p' l) [ #_ #E normalize in E; destruct ]
    * #a' #l' #H #E whd in E:(??%?); destruct
    cases (H a l' (refl ??)) * * #L1 #L2 #L3 #CARD
    @conj try @conj try @conj
    [ @L1 | @L2 | * /2/ #q cases (L3 q) [ /2/ | #E %2 @E ] | cases x normalize >CARD // ]
  ]
] qed.

lemma pm_try_remove_some' : ∀A,p,t,a.
  lookup_opt A p t = Some A a →
  ∃t'. pm_try_remove A p t = Some ? 〈a,t'〉.
#A #p elim p
[ * [ #a #L normalize in L; destruct
    | #q #l #r #a #L normalize in L; destruct % [2: % | skip ]
    ]
| #q #IH *
  [ #a #L normalize in L; destruct
  | #x #l #r #a #L cases (IH r a L) #r' #R
    % [2: whd in ⊢ (??%?); >R in ⊢ (??%?); % | skip ]
  ]
| #q #IH *
  [ #a #L normalize in L; destruct
  | #x #l #r #a #L cases (IH l a L) #l' #R
    % [2: whd in ⊢ (??%?); >R in ⊢ (??%?); % | skip ]
  ]
] qed.

(* An "informative" dependently typed fold *)

let rec pm_fold_inf_aux
 (A, B: Type[0])
 (t: positive_map A)
 (f: ∀p:Pos. ∀a:A. lookup_opt … p t = Some A a → B → B)
 (t': positive_map A)
 (pre: Pos → Pos)
 (b: B) on t': (∀p. lookup_opt … p t' = lookup_opt … (pre p) t) → B ≝
match t' return λt'. (∀p. lookup_opt A p t' = lookup_opt A (pre p) t) → B with
[ pm_leaf ⇒ λ_. b
| pm_node a l r ⇒ λH.
    let b ≝ match a return λa. lookup_opt A one (pm_node A a ??) = ? → B with [ None ⇒ λ_.b | Some a ⇒ λH. f (pre one) a ? b ] (H one) in
    let b ≝ pm_fold_inf_aux A B t f l (λx. pre (p0 x)) b ? in
            pm_fold_inf_aux A B t f r (λx. pre (p1 x)) b ?
].
[ #p @(H (p1 p)) | #p @(H (p0 p)) | <H % ]
qed.

definition pm_fold_inf : ∀A, B: Type[0]. ∀t: positive_map A.
 ∀f: (∀p:Pos. ∀a:A. lookup_opt … p t = Some A a → B → B). B → B ≝
  λA,B,t,f,b. pm_fold_inf_aux A B t f t (λx.x) b (λp. refl ??).

(* Find an entry in the map matching the given predicate *)

let rec pm_find_aux (pre: Pos → Pos) (A:Type[0]) (t: positive_map A) (p:Pos → A → bool) on t : option (Pos × A) ≝
match t with
[ pm_leaf ⇒ None ?
| pm_node a l r ⇒
  let x ≝ match a with
  [ Some a ⇒ if p (pre one) a then Some ? 〈pre one, a〉 else None ?
  | None ⇒ None ?
  ] in
  match x with
  [ Some x ⇒ Some ? x
  | None ⇒ 
    match pm_find_aux (λx. pre (p0 x)) A l p with
    [ None ⇒ pm_find_aux (λx. pre (p1 x)) A r p
    | Some y ⇒ Some ? y
    ]
  ]
].

definition pm_find : ∀A:Type[0]. positive_map A → (Pos → A → bool) → option (Pos × A) ≝
  pm_find_aux (λx.x).

lemma pm_find_aux_pre : ∀A,t,pre,p,q,a.
  pm_find_aux pre A t p = Some ? 〈q,a〉 →
  ∃q'. q = pre q'.
#A #t elim t
[ normalize #pre #p #q #a #E destruct
| * [2:#a] #l #r #IHl #IHr #pre #p #q #a' normalize
  [ cases (p (pre one) a) normalize [ #E destruct /2/ ] ]
  lapply (IHl (λx. pre (p0 x)) p)
  cases (pm_find_aux ?? l ?)
  [ 1,3: #_ normalize
    lapply (IHr (λx. pre (p1 x)) p)
    cases (pm_find_aux ?? r ?)
    [ 1,3: #_ #E destruct
    | *: * #q' #a'' #H #E destruct cases (H ?? (refl ??)) #q'' #E destruct /2/
    ]
  | *: * #q' #a'' #H normalize #E destruct cases (H ?? (refl ??)) #q'' #E destruct /2/
  ]
] qed.

(* XXX: Hmm... destruct doesn't work properly with this, not sure why *)
lemma option_pair_f_eq : ∀A,B:Type[0].∀a,a',b,b'. ∀f:A → A.
  (∀a,a'. f a = f a' → a = a') →
  Some (A×B) 〈f a,b〉 = Some ? 〈f a',b'〉 →
  a = a' ∧ b = b'.
#A #B #a #a' #b #b' #f #EXT #E
cut (f a = f a') [ destruct (E) destruct (e0) @e2 (* WTF? *) ]
#E' >(EXT … E') in E ⊢ %; #E destruct /2/
qed.

lemma pm_find_lookup : ∀A,p,q,a,t.
  pm_find A t p = Some ? 〈q,a〉 →
  lookup_opt A q t = Some ? a.
#A #p #q #a normalize #t
change with ((λx.x) q) in match q in ⊢ (% → ?);
cut (∀y,z:Pos. (λx.x) y = (λx.x) z → y = z) //
generalize in match q; -q generalize in match (λx.x);
elim t
[ #pre #q #_  normalize #E destruct
| * [2:#a'] #l #r #IHl #IHr #pre #q #PRE
  normalize [ cases (p (pre one) a') normalize [ #E cases (option_pair_f_eq … pre PRE E) #E1 #E2 destruct normalize % ] ]
  lapply (IHl (λx.pre (p0 x)))
  lapply (pm_find_aux_pre A l (λx.pre (p0 x)) p)
  cases (pm_find_aux ?? l ?)
  [ 1,3: #_ #_ normalize
    lapply (IHr (λx.pre (p1 x)))
    lapply (pm_find_aux_pre A r (λx.pre (p1 x)) p)
    cases (pm_find_aux ?? r ?)
    [ 1,3: #_ #_ #E destruct
    | *: * #q' #a' #H1 #H2 #E destruct cases (H1 ?? (refl ??)) #q'' #E lapply (PRE … E) #E' destruct
         whd in ⊢ (??%?); @H2 //
         #y #z #E lapply (PRE … E) #E' destruct //
    ]
  | *: * #q' #a' #H1 #H2 normalize #E destruct cases (H1 ?? (refl ??)) #q'' #E lapply (PRE … E) #E' destruct
         whd in ⊢ (??%?); @H2 //
         #y #z #E lapply (PRE … E) #E' destruct //
  ]
] qed.

lemma pm_find_aux_none : ∀A,t,pre,p,q,a.
  pm_find_aux pre A t p = None ? →
  lookup_opt A q t = Some ? a →
  ¬ p (pre q) a.
#A #t elim t
[ #pre #p #q #a #_ normalize #E destruct
| #oa #l #r #IHl #IHr #pre #p *
  [ #a cases oa
    [ #_ normalize #E destruct
    | #a' #F whd in F:(??%?); whd in F:(??match % with [_⇒?|_⇒?]?);
      #E normalize in E; destruct
      cases (p (pre one) a) in F ⊢ %; //
      normalize #E destruct
    ]
  | #q #a #F #L @(IHr (λx. pre (p1 x)) p q a ? L)
    whd in F:(??%?); cases oa in F;
    [2: #a' normalize cases (p (pre one) a') [ normalize #E destruct ] ]
    normalize cases (pm_find_aux ?? l p) normalize //
    #x #E destruct
  | #q #a #F #L @(IHl (λx. pre (p0 x)) p q a ? L)
    whd in F:(??%?); cases oa in F;
    [2: #a' normalize cases (p (pre one) a') [ normalize #E destruct ] ]
    normalize cases (pm_find_aux ?? l p) normalize //
  ]
] qed.

lemma pm_find_none : ∀A,t,p,q,a.
  pm_find A t p = None ? →
  lookup_opt A q t = Some ? a →
  ¬ p q a.
#A #t
@(pm_find_aux_none A t (λx.x))
qed.

lemma pm_find_predicate : ∀A,t,p,q,a.
  pm_find A t p = Some ? 〈q,a〉 →
  p q a.
#A #t #p normalize #q
change with ((λx.x) q) in match q; generalize in match q; -q generalize in match (λx.x);
elim t
[ normalize #pre #q #a #E destruct
| * [2:#a'] #l #r #IHl #IHr #pre #q #a normalize
  [ lapply (refl ? (p (pre one) a')) cases (p (pre one) a') in ⊢ (???% → %);
    [ #E normalize #E'
      cut (pre one = pre q) [ destruct (E') destruct (e0) @e2 (* XXX ! *) ]
      #E'' <E'' in E' ⊢ %; #E' destruct >E %
    | #_ normalize
    ]
  ]
  lapply (IHl (λx.pre (p0 x)))
  lapply (pm_find_aux_pre A l (λx.pre (p0 x)) p)
  cases (pm_find_aux ?? l ?)
  [ 1,3: #_ #_ normalize
    lapply (IHr (λx.pre (p1 x)))
    lapply (pm_find_aux_pre A r (λx.pre (p1 x)) p)
    cases (pm_find_aux ?? r ?)
    [ 1,3: #_ #_ #E destruct
    | *: * #q' #a' #H1 #H2 #E destruct cases (H1 … (refl ??)) #q'' #E >E @H2 >E %
    ]
  | *: * #q' #a' #H1 #H2 normalize #E destruct cases (H1 … (refl ??)) #q'' #E >E @H2 >E %
  ]
] qed.