(************************************************************************)
(*         *      The Rocq Prover / The Rocq Development Team           *)
(*  v      *         Copyright INRIA, CNRS and contributors             *)
(* <O___,, * (see version control and CREDITS file for authors & dates) *)
(*   \VV/  **************************************************************)
(*    //   *    This file is distributed under the terms of the         *)
(*         *     GNU Lesser General Public License Version 2.1          *)
(*         *     (see LICENSE file for the text of the license)         *)
(************************************************************************)
(* Finite map library.  *)
(** * FMapAVL *)
(** This module implements maps using AVL trees. It follows the implementation from Ocaml's standard library.  See the comments at the beginning of FSetAVL for more details. *)
Set Implicit Arguments.
Unset Strict Implicit.
(** Notations and helper lemma about pairs *)
Declare Scope pair_scope.
Notation "s #1" := (fst s) (at level 9, format "s '#1'") : pair_scope.
Notation "s #2" := (snd s) (at level 9, format "s '#2'") : pair_scope.
(** * The Raw functor  Functor of pure functions + separate proofs of invariant preservation *)
Module Raw (Import I:Int)(X: OrderedType).
Local Open Scope pair_scope.
Local Open Scope lazy_bool_scope.
Local Open Scope Int_scope.
Local Notation int := I.t.
Definition key := X.t.
#[global]
Hint Transparent key : core.
(** * Trees *)
(** * Trees  The fifth field of [Node] is the height of the tree *)
#[universes(template)]
Inductive tree {elt : Type} :=
| Leaf : tree
| Node : tree -> key -> elt -> tree -> int -> tree.
Arguments tree : clear implicits.
Section Elt.
Variable elt : Type.
Notation t := (tree elt).
Implicit Types m : t.
(** * Basic functions on trees: height and cardinal *)
Definition height (m : t) : int :=
match m with
| Leaf => 0
| Node _ _ _ _ h => h
end.
Fixpoint cardinal (m : t) : nat :=
match m with
| Leaf => 0%nat
| Node l _ _ r _ => S (cardinal l + cardinal r)
end.
(** * Empty Map *)
Definition empty : t := Leaf.
(** * Emptyness test *)
Definition is_empty m := match m with Leaf => true | _ => false end.
(** * Membership *)
(** The [mem] function is deciding membership. It exploits the [bst] property to achieve logarithmic complexity. *)
Fixpoint mem x m : bool :=
match m with
|  Leaf => false
|  Node l y _ r _ => match X.compare x y with
| LT _ => mem x l
| EQ _ => true
| GT _ => mem x r
end
end.
Fixpoint find x m : option elt :=
match m with
|  Leaf => None
|  Node l y d r _ => match X.compare x y with
| LT _ => find x l
| EQ _ => Some d
| GT _ => find x r
end
end.
(** * Helper functions *)
(** [create l x r] creates a node, assuming [l] and [r] to be balanced and [|height l - height r| <= 2]. *)
Definition create l x e r :=
Node l x e r (max (height l) (height r) + 1).
(** [bal l x e r] acts as [create], but performs one step of rebalancing if necessary, i.e. assumes [|height l - height r| <= 3]. *)
Definition assert_false := create.
Definition bal l x d r :=
let hl := height l in
let hr := height r in
if gt_le_dec hl (hr+2) then
match l with
| Leaf => assert_false l x d r
| Node ll lx ld lr _ =>
if ge_lt_dec (height ll) (height lr) then
create ll lx ld (create lr x d r)
else
match lr with
| Leaf => assert_false l x d r
| Node lrl lrx lrd lrr _ =>
create (create ll lx ld lrl) lrx lrd (create lrr x d r)
end
end
else
if gt_le_dec hr (hl+2) then
match r with
| Leaf => assert_false l x d r
| Node rl rx rd rr _ =>
if ge_lt_dec (height rr) (height rl) then
create (create l x d rl) rx rd rr
else
match rl with
| Leaf => assert_false l x d r
| Node rll rlx rld rlr _ =>
create (create l x d rll) rlx rld (create rlr rx rd rr)
end
end
else
create l x d r.
(** * Insertion *)
Fixpoint add x d m :=
match m with
| Leaf => Node Leaf x d Leaf 1
| Node l y d' r h =>
match X.compare x y with
| LT _ => bal (add x d l) y d' r
| EQ _ => Node l y d r h
| GT _ => bal l y d' (add x d r)
end
end.
(** * Extraction of minimum binding  Morally, [remove_min] is to be applied to a non-empty tree [t = Node l x e r h]. Since we can't deal here with [assert false] for [t=Leaf], we pre-unpack [t] (and forget about [h]). *)
Fixpoint remove_min l x d r : t*(key*elt) :=
match l with
| Leaf => (r,(x,d))
| Node ll lx ld lr lh =>
let (l',m) := remove_min ll lx ld lr in
(bal l' x d r, m)
end.
(** * Merging two trees  [merge t1 t2] builds the union of [t1] and [t2] assuming all elements of [t1] to be smaller than all elements of [t2], and [|height t1 - height t2| <= 2]. *)
Definition merge s1 s2 :=  match s1,s2 with
| Leaf, _ => s2
| _, Leaf => s1
| _, Node l2 x2 d2 r2 h2 =>
match remove_min l2 x2 d2 r2 with
(s2',(x,d)) => bal s1 x d s2'
end
end.
(** * Deletion *)
Fixpoint remove x m := match m with
| Leaf => Leaf
| Node l y d r h =>
match X.compare x y with
| LT _ => bal (remove x l) y d r
| EQ _ => merge l r
| GT _ => bal l y d (remove x r)
end
end.
(** * join  Same as [bal] but does not assume anything regarding heights of [l] and [r]. *)
Fixpoint join l : key -> elt -> t -> t :=
match l with
| Leaf => add
| Node ll lx ld lr lh => fun x d =>
fix join_aux (r:t) : t := match r with
| Leaf =>  add x d l
| Node rl rx rd rr rh =>
if gt_le_dec lh (rh+2) then bal ll lx ld (join lr x d r)
else if gt_le_dec rh (lh+2) then bal (join_aux rl) rx rd rr
else create l x d r
end
end.
(** * Splitting  [split x m] returns a triple [(l, o, r)] where - [l] is the set of elements of [m] that are [< x] - [r] is the set of elements of [m] that are [> x] - [o] is the result of [find x m]. *)
Record triple := mktriple { t_left:t; t_opt:option elt; t_right:t }.
Notation "<< l , b , r >>" := (mktriple l b r) (at level 9).
Fixpoint split x m : triple := match m with
| Leaf => << Leaf, None, Leaf >>
| Node l y d r h =>
match X.compare x y with
| LT _ => let (ll,o,rl) := split x l in << ll, o, join rl y d r >>
| EQ _ => << l, Some d, r >>
| GT _ => let (rl,o,rr) := split x r in << join l y d rl, o, rr >>
end
end.
(** * Concatenation  Same as [merge] but does not assume anything about heights. *)
Definition concat m1 m2 :=
match m1, m2 with
| Leaf, _ => m2
| _ , Leaf => m1
| _, Node l2 x2 d2 r2 _ =>
let (m2',xd) := remove_min l2 x2 d2 r2 in
join m1 xd#1 xd#2 m2'
end.
(** * Elements *)
(** [elements_tree_aux acc t] catenates the elements of [t] in infix order to the list [acc] *)
Fixpoint elements_aux (acc : list (key*elt)) m : list (key*elt) :=
match m with
| Leaf => acc
| Node l x d r _ => elements_aux ((x,d) :: elements_aux acc r) l
end.
(** then [elements] is an instantiation with an empty [acc] *)
Definition elements := elements_aux nil.
(** * Fold *)
Fixpoint fold (A : Type) (f : key -> elt -> A -> A) (m : t) : A -> A :=
fun a => match m with
| Leaf => a
| Node l x d r _ => fold f r (f x d (fold f l a))
end.
(** * Comparison *)
Variable cmp : elt->elt->bool.
(** ** Enumeration of the elements of a tree *)
Inductive enumeration :=
| End : enumeration
| More : key -> elt -> t -> enumeration -> enumeration.
(** [cons m e] adds the elements of tree [m] on the head of enumeration [e]. *)
Fixpoint cons m e : enumeration :=
match m with
| Leaf => e
| Node l x d r h => cons l (More x d r e)
end.
(** One step of comparison of elements *)
Definition equal_more x1 d1 (cont:enumeration->bool) e2 :=
match e2 with
| End => false
| More x2 d2 r2 e2 =>
match X.compare x1 x2 with
| EQ _ => cmp d1 d2 &&& cont (cons r2 e2)
| _ => false
end
end.
(** Comparison of left tree, middle element, then right tree *)
Fixpoint equal_cont m1 (cont:enumeration->bool) e2 :=
match m1 with
| Leaf => cont e2
| Node l1 x1 d1 r1 _ =>
equal_cont l1 (equal_more x1 d1 (equal_cont r1 cont)) e2
end.
(** Initial continuation *)
Definition equal_end e2 := match e2 with End => true | _ => false end.
(** The complete comparison *)
Definition equal m1 m2 := equal_cont m1 equal_end (cons m2 End).
End Elt.
Notation t := tree.
Arguments Leaf : clear implicits.
Arguments Node [elt].
Notation "<< l , b , r >>" := (mktriple l b r) (at level 9).
Notation "t #l" := (t_left t) (at level 9, format "t '#l'").
Notation "t #o" := (t_opt t) (at level 9, format "t '#o'").
Notation "t #r" := (t_right t) (at level 9, format "t '#r'").
(** * Map *)
Fixpoint map (elt elt' : Type)(f : elt -> elt')(m : t elt) : t elt' :=
match m with
| Leaf _   => Leaf _
| Node l x d r h => Node (map f l) x (f d) (map f r) h
end.
(* * Mapi *)
Fixpoint mapi (elt elt' : Type)(f : key -> elt -> elt')(m : t elt) : t elt' :=
match m with
| Leaf _ => Leaf _
| Node l x d r h => Node (mapi f l) x (f x d) (mapi f r) h
end.
(** * Map with removal *)
Fixpoint map_option (elt elt' : Type)(f : key -> elt -> option elt')(m : t elt)
: t elt' :=
match m with
| Leaf _ => Leaf _
| Node l x d r h =>
match f x d with
| Some d' => join (map_option f l) x d' (map_option f r)
| None => concat (map_option f l) (map_option f r)
end
end.
(** * Optimized map2  Suggestion by B. Gregoire: a [map2] function with specialized arguments that allows bypassing some tree traversal. Instead of one [f0] of type [key -> option elt -> option elt' -> option elt''], we ask here for: - [f] which is a specialisation of [f0] when first option isn't [None] - [mapl] treats a [tree elt] with [f0] when second option is [None] - [mapr] treats a [tree elt'] with [f0] when first option is [None]  The idea is that [mapl] and [mapr] can be instantaneous (e.g. the identity or some constant function). *)
Section Map2_opt.
Variable elt elt' elt'' : Type.
Variable f : key -> elt -> option elt' -> option elt''.
Variable mapl : t elt -> t elt''.
Variable mapr : t elt' -> t elt''.
Fixpoint map2_opt m1 m2 :=
match m1, m2 with
| Leaf _, _ => mapr m2
| _, Leaf _ => mapl m1
| Node l1 x1 d1 r1 h1, _ =>
let (l2',o2,r2') := split x1 m2 in
match f x1 d1 o2 with
| Some e => join (map2_opt l1 l2') x1 e (map2_opt r1 r2')
| None => concat (map2_opt l1 l2') (map2_opt r1 r2')
end
end.
End Map2_opt.
(** * Map2  The [map2] function of the Map interface can be implemented via [map2_opt] and [map_option]. *)
Section Map2.
Variable elt elt' elt'' : Type.
Variable f : option elt -> option elt' -> option elt''.
Definition map2 : t elt -> t elt' -> t elt'' :=
map2_opt
(fun _ d o => f (Some d) o)
(map_option (fun _ d => f (Some d) None))
(map_option (fun _ d' => f None (Some d'))).
End Map2.
(** * Invariants *)
Section Invariants.
Variable elt : Type.
(** ** Occurrence in a tree *)
Inductive MapsTo (x : key)(e : elt) : t elt -> Prop :=
| MapsRoot : forall l r h y,
X.eq x y -> MapsTo x e (Node l y e r h)
| MapsLeft : forall l r h y e',
MapsTo x e l -> MapsTo x e (Node l y e' r h)
| MapsRight : forall l r h y e',
MapsTo x e r -> MapsTo x e (Node l y e' r h).
Inductive In (x : key) : t elt -> Prop :=
| InRoot : forall l r h y e,
X.eq x y -> In x (Node l y e r h)
| InLeft : forall l r h y e',
In x l -> In x (Node l y e' r h)
| InRight : forall l r h y e',
In x r -> In x (Node l y e' r h).
Definition In0 k m := exists e:elt, MapsTo k e m.
(** ** Binary search trees *)
(** [lt_tree x s]: all elements in [s] are smaller than [x] (resp. greater for [gt_tree]) *)
Definition lt_tree x m := forall y, In y m -> X.lt y x.
Definition gt_tree x m := forall y, In y m -> X.lt x y.
(** [bst t] : [t] is a binary search tree *)
Inductive bst : t elt -> Prop :=
| BSLeaf : bst (Leaf _)
| BSNode : forall x e l r h, bst l -> bst r ->
lt_tree x l -> gt_tree x r -> bst (Node l x e r h).
End Invariants.
(** * Correctness proofs, isolated in a sub-module *)
Module Proofs.
Module MX := OrderedTypeFacts X.
Module PX := KeyOrderedType X.
Module L := FMapList.Raw X.
#[local] Ltac caseq :=
match goal with [ |- context [match ?t with _ => _ end] ] =>
let cmp := fresh in
let H := fresh in
remember t as cmp eqn:H; symmetry in H; destruct cmp
end.

Lemma map_option_ind [elt elt' : Type] [f : key -> elt -> option elt'] [P : t elt -> t elt' -> Prop] :
(forall m : t elt, m = Leaf elt -> P (Leaf elt) (Leaf elt')) ->
(forall (m l : t elt) (x : key) (d : elt) (r : t elt) (_x : int),
m = Node l x d r _x ->
forall d' : elt',
f x d = Some d' ->
P l (map_option f l) ->
P r (map_option f r) -> P (Node l x d r _x) (join (map_option f l) x d' (map_option f r))) ->
(forall (m l : t elt) (x : key) (d : elt) (r : t elt) (_x : int),
m = Node l x d r _x ->
f x d = None ->
P l (map_option f l) ->
P r (map_option f r) -> P (Node l x d r _x) (concat (map_option f l) (map_option f r))) ->
forall m : t elt, P m (map_option f m).