include "basics/types.ma".
include "utilities/binary/positive.ma".

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 insert (A:Type[0]) (b:Pos) (a:A) (t:positive_map A) on b : positive_map A ≝
match b with
[ one ⇒
    match t with
    [ pm_leaf ⇒ pm_node A (Some ? a) (pm_leaf A) (pm_leaf A)
    | pm_node _ l r ⇒ pm_node A (Some ? a) l r
    ]
| p0 tl ⇒
    match t with
    [ pm_leaf ⇒ pm_node A (None ?) (insert A tl a (pm_leaf A)) (pm_leaf A)
    | pm_node x l r ⇒ pm_node A x (insert A tl a l) r
    ]
| p1 tl ⇒
    match t with
    [ pm_leaf ⇒ pm_node A (None ?) (pm_leaf A) (insert A tl a (pm_leaf A))
    | pm_node x l r ⇒ pm_node A x l (insert A tl a r)
    ]
].

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_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.

axiom lookup_insert_hit:
  ∀a: Type[0].
  ∀v: a.
  ∀b: Pos.
  ∀t: positive_map a.
  ∀d: a.
    lookup … b (insert … b v t) d = v.

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
].    

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.

let rec merge (A: Type[0]) (b: positive_map A) (c:positive_map A) on b: positive_map A ≝
match b with
[ pm_leaf ⇒ c
| pm_node a l r ⇒
  match c with
  [ pm_leaf ⇒ pm_node A a l r
  | pm_node a' l' r' ⇒ pm_node A a' (merge A l l') (merge A r r')
  ]
].