(* -*- coding: utf-8 -*- *)
(************************************************************************)
(*         *      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)         *)
(************************************************************************)
(** * MSetAVL : Implementation of MSetInterface via AVL trees *)
(** This module implements finite sets using AVL trees. It follows the implementation from Ocaml's standard library,  All operations given here expect and produce well-balanced trees (in the ocaml sense: heights of subtrees shouldn't differ by more than 2), and hence has low complexities (e.g. add is logarithmic in the size of the set). But proving these balancing preservations is in fact not necessary for ensuring correct operational behavior and hence fulfilling the MSet interface. As a consequence, balancing results are not part of this file anymore, they can now be found in [MSetFullAVL].  Four operations ([union], [subset], [compare] and [equal]) have been slightly adapted in order to have only structural recursive calls. The precise ocaml versions of these operations have also been formalized (thanks to Function+measure), see [ocaml_union], [ocaml_subset], [ocaml_compare] and [ocaml_equal] in [MSetFullAVL]. The structural variants compute faster in Coq, whereas the other variants produce nicer and/or (slightly) faster code after extraction. *)
Set Implicit Arguments.
Unset Strict Implicit.
(* for nicer extraction, we create inductive principles only when needed *)
Local Unset Elimination Schemes.
(** * Ops : the pure functions *)
Module Ops (Import I:Int)(X:OrderedType) <: MSetInterface.Ops X.
Local Open Scope Int_scope.
Local Notation int := I.t.
(** ** Generic trees instantiated with integer height *)
(** We reuse a generic definition of trees where the information parameter is a [Int.t]. Functions like mem or fold are also provided by this generic functor. *)
Include MSetGenTree.Ops X I.
Definition t := tree.
(** ** Height of trees *)
Definition height (s : t) : int :=
match s with
| Leaf => 0
| Node h _ _ _ => h
end.
(** ** Singleton set *)
Definition singleton x := Node 1 Leaf x Leaf.
(** ** 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 r :=
Node (max (height l) (height r) + 1) l x r.
(** [bal l x 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 r :=
let hl := height l in
let hr := height r in
if (hr+2) <? hl then
match l with
| Leaf => assert_false l x r
| Node _ ll lx lr =>
if (height lr) <=? (height ll) then
create ll lx (create lr x r)
else
match lr with
| Leaf => assert_false l x r
| Node _ lrl lrx lrr =>
create (create ll lx lrl) lrx (create lrr x r)
end
end
else
if (hl+2) <? hr then
match r with
| Leaf => assert_false l x r
| Node _ rl rx rr =>
if (height rl) <=? (height rr) then
create (create l x rl) rx rr
else
match rl with
| Leaf => assert_false l x r
| Node _ rll rlx rlr =>
create (create l x rll) rlx (create rlr rx rr)
end
end
else
create l x r.
(** ** Insertion *)
Fixpoint add x s := match s with
| Leaf => Node 1 Leaf x Leaf
| Node h l y r =>
match X.compare x y with
| Lt => bal (add x l) y r
| Eq => Node h l y r
| Gt => bal l y (add x r)
end
end.
(** ** Join  Same as [bal] but does not assume anything regarding heights of [l] and [r]. *)
Fixpoint join l : elt -> t -> t :=
match l with
| Leaf => add
| Node lh ll lx lr => fun x =>
fix join_aux (r:t) : t := match r with
| Leaf => add x l
| Node rh rl rx rr =>
if (rh+2) <? lh then bal ll lx (join lr x r)
else if (lh+2) <? rh then bal (join_aux rl) rx rr
else create l x r
end
end.
(** ** Extraction of minimum element  Morally, [remove_min] is to be applied to a non-empty tree [t = Node h l x r]. 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 r : t*elt :=
match l with
| Leaf => (r,x)
| Node lh ll lx lr =>
let (l',m) := remove_min ll lx lr in (bal l' x 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 r2 =>
let (s2',m) := remove_min l2 x2 r2 in bal s1 m s2'
end.
(** ** Deletion *)
Fixpoint remove x s := match s with
| Leaf => Leaf
| Node _ l y r =>
match X.compare x y with
| Lt => bal (remove x l) y r
| Eq => merge l r
| Gt => bal l y (remove x r)
end
end.
(** ** Concatenation  Same as [merge] but does not assume anything about heights. *)
Definition concat s1 s2 :=
match s1, s2 with
| Leaf, _ => s2
| _, Leaf => s1
| _, Node _ l2 x2 r2 =>
let (s2',m) := remove_min l2 x2 r2 in
join s1 m s2'
end.
(** ** Splitting  [split x s] returns a triple [(l, present, r)] where - [l] is the set of elements of [s] that are [< x] - [r] is the set of elements of [s] that are [> x] - [present] is [true] if and only if [s] contains  [x]. *)
Record triple := mktriple { t_left:t; t_in:bool; t_right:t }.
Notation "<< l , b , r >>" := (mktriple l b r) (at level 9).
Fixpoint split x s : triple := match s with
| Leaf => << Leaf, false, Leaf >>
| Node _ l y r =>
match X.compare x y with
| Lt => let (ll,b,rl) := split x l in << ll, b, join rl y r >>
| Eq => << l, true, r >>
| Gt => let (rl,b,rr) := split x r in << join l y rl, b, rr >>
end
end.
(** ** Intersection *)
Fixpoint inter s1 s2 := match s1, s2 with
| Leaf, _ => Leaf
| _, Leaf => Leaf
| Node _ l1 x1 r1, _ =>
let (l2',pres,r2') := split x1 s2 in
if pres then join (inter l1 l2') x1 (inter r1 r2')
else concat (inter l1 l2') (inter r1 r2')
end.
(** ** Difference *)
Fixpoint diff s1 s2 := match s1, s2 with
| Leaf, _ => Leaf
| _, Leaf => s1
| Node _ l1 x1 r1, _ =>
let (l2',pres,r2') := split x1 s2 in
if pres then concat (diff l1 l2') (diff r1 r2')
else join (diff l1 l2') x1 (diff r1 r2')
end.
(** ** Union *)
(** In ocaml, heights of [s1] and [s2] are compared each time in order to recursively perform the split on the smaller set. Unfortunately, this leads to a non-structural algorithm. The following code is a simplification of the ocaml version: no comparison of heights. It might be slightly slower, but experimentally all the tests I've made in ocaml have shown this potential slowdown to be non-significant. Anyway, the exact code of ocaml has also been formalized thanks to Function+measure, see [ocaml_union] in [MSetFullAVL]. *)
Fixpoint union s1 s2 :=
match s1, s2 with
| Leaf, _ => s2
| _, Leaf => s1
| Node _ l1 x1 r1, _ =>
let (l2',_,r2') := split x1 s2 in
join (union l1 l2') x1 (union r1 r2')
end.
(** ** Filter *)
Fixpoint filter (f:elt->bool) s := match s with
| Leaf => Leaf
| Node _ l x r =>
let l' := filter f l in
let r' := filter f r in
if f x then join l' x r' else concat l' r'
end.
(** ** Partition *)
Fixpoint partition (f:elt->bool)(s : t) : t*t :=
match s with
| Leaf => (Leaf, Leaf)
| Node _ l x r =>
let (l1,l2) := partition f l in
let (r1,r2) := partition f r in
if f x then (join l1 x r1, concat l2 r2)
else (concat l1 r1, join l2 x r2)
end.
End Ops.
(** * MakeRaw  Functor of pure functions + a posteriori proofs of invariant preservation *)
Module MakeRaw (Import I:Int)(X:OrderedType) <: RawSets X.
Include Ops I X.
(** Generic definition of binary-search-trees and proofs of specifications for generic functions such as mem or fold. *)
Include MSetGenTree.Props X I.
(** Automation and dedicated tactics *)
Local Hint Immediate MX.eq_sym : core.
Local Hint Unfold In lt_tree gt_tree Ok : core.
Local Hint Constructors InT bst : core.
Local Hint Resolve MX.eq_refl MX.eq_trans MX.lt_trans ok : core.
Local Hint Resolve lt_leaf gt_leaf lt_tree_node gt_tree_node : core.
Local Hint Resolve lt_tree_not_in lt_tree_trans gt_tree_not_in gt_tree_trans : core.
Local Hint Resolve elements_spec2 : core.
(* Sometimes functional induction will expose too much of a tree structure. The following tactic allows factoring back a Node whose internal parts occurs nowhere else. *)
(* TODO: why Ltac instead of Tactic Notation don't work ? why clear ? *)
Tactic Notation "factornode" ident(s) :=
try clear s;
match goal with
| |- context [Node ?l ?x ?r ?h] =>
set (s:=Node l x r h) in *; clearbody s; clear l x r h
| _ : context [Node ?l ?x ?r ?h] |- _ =>
set (s:=Node l x r h) in *; clearbody s; clear l x r h
end.
(** Inductions principles for some of the set operators *)
#[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 union_ind [P : tree -> tree -> tree -> Prop] :
(forall s1 s2 : tree, s1 = Leaf -> P Leaf s2 s2) ->
(forall (s1 s2 : tree) (_x : I.t) (l1 : tree) (x1 : X.t) (r1 : tree),
s1 = Node _x l1 x1 r1 -> s2 = Leaf -> P (Node _x l1 x1 r1) Leaf s1) ->
(forall (s1 s2 : tree) (_x : I.t) (l1 : tree) (x1 : X.t) (r1 : tree),
s1 = Node _x l1 x1 r1 ->
forall (_x0 : I.t) (_x1 : tree) (_x2 : X.t) (_x3 : tree),
s2 = Node _x0 _x1 _x2 _x3 ->
forall (l2' : t) (_x4 : bool) (r2' : t),
split x1 s2 = << l2', _x4, r2' >> ->
P l1 l2' (union l1 l2') ->
P r1 r2' (union r1 r2') ->
P (Node _x l1 x1 r1) (Node _x0 _x1 _x2 _x3) (join (union l1 l2') x1 (union r1 r2'))) ->
forall s1 s2 : tree, P s1 s2 (union s1 s2).