(** The type of color annotation. *)
Inductive color := Red | Black.
Module Color.
Definition t := color.
End Color.
(** * Ops : the pure functions *)
Module Ops (X:Orders.OrderedType) <: MSetInterface.Ops X.
(** ** Generic trees instantiated with color *)
(** We reuse a generic definition of trees where the information parameter is a color. Functions like mem or fold are also provided by this generic functor. *)
Include MSetGenTree.Ops X Color.
Definition t := tree.
Local Notation Rd := (Node Red).
Local Notation Bk := (Node Black).
(** ** Basic tree *)
Definition singleton (k: elt) : tree := Bk Leaf k Leaf.
(** ** Changing root color *)
Definition makeBlack t :=
match t with
| Leaf => Leaf
| Node _ a x b => Bk a x b
end.
Definition makeRed t :=
match t with
| Leaf => Leaf
| Node _ a x b => Rd a x b
end.
(** ** Balancing *)
(** We adapt when one side is not a true red-black tree. Both sides have the same black depth. *)
Definition lbal l k r :=
match l with
| Rd (Rd a x b) y c => Rd (Bk a x b) y (Bk c k r)
| Rd a x (Rd b y c) => Rd (Bk a x b) y (Bk c k r)
| _ => Bk l k r
end.
Definition rbal l k r :=
match r with
| Rd (Rd b y c) z d => Rd (Bk l k b) y (Bk c z d)
| Rd b y (Rd c z d) => Rd (Bk l k b) y (Bk c z d)
| _ => Bk l k r
end.
(** A variant of [rbal], with reverse pattern order. Is it really useful ? Should we always use it ? *)
Definition rbal' l k r :=
match r with
| Rd b y (Rd c z d) => Rd (Bk l k b) y (Bk c z d)
| Rd (Rd b y c) z d => Rd (Bk l k b) y (Bk c z d)
| _ => Bk l k r
end.
(** Balancing with different black depth. One side is almost a red-black tree, while the other is a true red-black tree, but with black depth + 1. Used in deletion. *)
Definition lbalS l k r :=
match l with
| Rd a x b => Rd (Bk a x b) k r
| _ =>
match r with
| Bk a y b => rbal' l k (Rd a y b)
| Rd (Bk a y b) z c => Rd (Bk l k a) y (rbal' b z (makeRed c))
| _ => Rd l k r (* impossible *)
end
end.
Definition rbalS l k r :=
match r with
| Rd b y c => Rd l k (Bk b y c)
| _ =>
match l with
| Bk a x b => lbal (Rd a x b) k r
| Rd a x (Bk b y c) => Rd (lbal (makeRed a) x b) y (Bk c k r)
| _ => Rd l k r (* impossible *)
end
end.
(** ** Insertion *)
Fixpoint ins x s :=
match s with
| Leaf => Rd Leaf x Leaf
| Node c l y r =>
match X.compare x y with
| Eq => s
| Lt =>
match c with
| Red => Rd (ins x l) y r
| Black => lbal (ins x l) y r
end
| Gt =>
match c with
| Red => Rd l y (ins x r)
| Black => rbal l y (ins x r)
end
end
end.
Definition add x s := makeBlack (ins x s).
(** ** Deletion *)
Fixpoint append (l:tree) : tree -> tree :=
match l with
| Leaf => fun r => r
| Node lc ll lx lr =>
fix append_l (r:tree) : tree :=
match r with
| Leaf => l
| Node rc rl rx rr =>
match lc, rc with
| Red, Red =>
let lrl := append lr rl in
match lrl with
| Rd lr' x rl' => Rd (Rd ll lx lr') x (Rd rl' rx rr)
| _ => Rd ll lx (Rd lrl rx rr)
end
| Black, Black =>
let lrl := append lr rl in
match lrl with
| Rd lr' x rl' => Rd (Bk ll lx lr') x (Bk rl' rx rr)
| _ => lbalS ll lx (Bk lrl rx rr)
end
| Black, Red => Rd (append_l rl) rx rr
| Red, Black => Rd ll lx (append lr r)
end
end
end.
Fixpoint del x t :=
match t with
| Leaf => Leaf
| Node _ a y b =>
match X.compare x y with
| Eq => append a b
| Lt =>
match a with
| Bk _ _ _ => lbalS (del x a) y b
| _ => Rd (del x a) y b
end
| Gt =>
match b with
| Bk _ _ _ => rbalS a y (del x b)
| _ => Rd a y (del x b)
end
end
end.
Definition remove x t := makeBlack (del x t).
(** ** Removing minimal element *)
Fixpoint delmin l x r : (elt * tree) :=
match l with
| Leaf => (x,r)
| Node lc ll lx lr =>
let (k,l') := delmin ll lx lr in
match lc with
| Black => (k, lbalS l' x r)
| Red => (k, Rd l' x r)
end
end.
Definition remove_min t : option (elt * tree) :=
match t with
| Leaf => None
| Node _ l x r =>
let (k,t) := delmin l x r in
Some (k, makeBlack t)
end.
(** ** Tree-ification  We rebuild a tree of size [if pred then n-1 else n] as soon as the list [l] has enough elements *)
Definition bogus : tree * list elt := (Leaf, nil).
Notation treeify_t := (list elt -> tree * list elt).
Definition treeify_zero : treeify_t :=
fun acc => (Leaf,acc).
Definition treeify_one : treeify_t :=
fun acc => match acc with
| x::acc => (Rd Leaf x Leaf, acc)
| _ => bogus
end.
Definition treeify_cont (f g : treeify_t) : treeify_t :=
fun acc =>
match f acc with
| (l, x::acc) =>
match g acc with
| (r, acc) => (Bk l x r, acc)
end
| _ => bogus
end.
Fixpoint treeify_aux (pred:bool)(n: positive) : treeify_t :=
match n with
| xH => if pred then treeify_zero else treeify_one
| xO n => treeify_cont (treeify_aux pred n) (treeify_aux true n)
| xI n => treeify_cont (treeify_aux false n) (treeify_aux pred n)
end.
Fixpoint plength_aux (l:list elt)(p:positive) := match l with
| nil => p
| _::l => plength_aux l (Pos.succ p)
end.
Definition plength l := plength_aux l 1.
Definition treeify (l:list elt) :=
fst (treeify_aux true (plength l) l).
(** ** Filtering *)
Fixpoint filter_aux (f: elt -> bool) s acc :=
match s with
| Leaf => acc
| Node _ l k r =>
let acc := filter_aux f r acc in
if f k then filter_aux f l (k::acc)
else filter_aux f l acc
end.
Definition filter (f: elt -> bool) (s: t) : t :=
treeify (filter_aux f s nil).
Fixpoint partition_aux (f: elt -> bool) s acc1 acc2 :=
match s with
| Leaf => (acc1,acc2)
| Node _ sl k sr =>
let (acc1, acc2) := partition_aux f sr acc1 acc2 in
if f k then partition_aux f sl (k::acc1) acc2
else partition_aux f sl acc1 (k::acc2)
end.
Definition partition (f: elt -> bool) (s:t) : t*t :=
let (ok,ko) := partition_aux f s nil nil in
(treeify ok, treeify ko).
(** ** Union, intersection, difference *)
(** union of the elements of [l1] and [l2] into a third [acc] list. *)
Fixpoint union_list l1 : list elt -> list elt -> list elt :=
match l1 with
| nil => @rev_append _
| x::l1' =>
fix union_l1 l2 acc :=
match l2 with
| nil => rev_append l1 acc
| y::l2' =>
match X.compare x y with
| Eq => union_list l1' l2' (x::acc)
| Lt => union_l1 l2' (y::acc)
| Gt => union_list l1' l2 (x::acc)
end
end
end.
Definition linear_union s1 s2 :=
treeify (union_list (rev_elements s1) (rev_elements s2) nil).
Fixpoint inter_list l1 : list elt -> list elt -> list elt :=
match l1 with
| nil => fun _ acc => acc
| x::l1' =>
fix inter_l1 l2 acc :=
match l2 with
| nil => acc
| y::l2' =>
match X.compare x y with
| Eq => inter_list l1' l2' (x::acc)
| Lt => inter_l1 l2' acc
| Gt => inter_list l1' l2 acc
end
end
end.
Definition linear_inter s1 s2 :=
treeify (inter_list (rev_elements s1) (rev_elements s2) nil).
Fixpoint diff_list l1 : list elt -> list elt -> list elt :=
match l1 with
| nil => fun _ acc => acc
| x::l1' =>
fix diff_l1 l2 acc :=
match l2 with
| nil => rev_append l1 acc
| y::l2' =>
match X.compare x y with
| Eq => diff_list l1' l2' acc
| Lt => diff_l1 l2' acc
| Gt => diff_list l1' l2 (x::acc)
end
end
end.
Definition linear_diff s1 s2 :=
treeify (diff_list (rev_elements s1) (rev_elements s2) nil).
(** [compare_height] returns: - [Lt] if [height s2] is at least twice [height s1]; - [Gt] if [height s1] is at least twice [height s2]; - [Eq] if heights are approximately equal. Warning: this is not an equivalence relation! but who cares.... *)
Definition skip_red t :=
match t with
| Rd t' _ _ => t'
| _ => t
end.
Definition skip_black t :=
match skip_red t with
| Bk t' _ _ => t'
| t' => t'
end.
Fixpoint compare_height (s1x s1 s2 s2x: tree) : comparison :=
match skip_red s1x, skip_red s1, skip_red s2, skip_red s2x with
| Node _ s1x' _ _, Node _ s1' _ _, Node _ s2' _ _, Node _ s2x' _ _ =>
compare_height (skip_black s1x') s1' s2' (skip_black s2x')
| _, Leaf, _, Node _ _ _ _ => Lt
| Node _ _ _ _, _, Leaf, _ => Gt
| Node _ s1x' _ _, Node _ s1' _ _, Node _ s2' _ _, Leaf =>
compare_height (skip_black s1x') s1' s2' Leaf
| Leaf, Node _ s1' _ _, Node _ s2' _ _, Node _ s2x' _ _ =>
compare_height Leaf s1'  s2'  (skip_black s2x')
| _, _, _, _ => Eq
end.
(** When one tree is quite smaller than the other, we simply adds repeatively all its elements in the big one. For trees of comparable height, we rather use [linear_union]. *)
Definition union (t1 t2: t) : t :=
match compare_height t1 t1 t2 t2 with
| Lt => fold add t1 t2
| Gt => fold add t2 t1
| Eq => linear_union t1 t2
end.
Definition diff (t1 t2: t) : t :=
match compare_height t1 t1 t2 t2 with
| Lt => filter (fun k => negb (mem k t2)) t1
| Gt => fold remove t2 t1
| Eq => linear_diff t1 t2
end.
Definition inter (t1 t2: t) : t :=
match compare_height t1 t1 t2 t2 with
| Lt => filter (fun k => mem k t2) t1
| Gt => filter (fun k => mem k t1) t2
| Eq => linear_inter t1 t2
end.
End Ops.
(** * MakeRaw : the pure functions and their specifications *)
Module Type MakeRaw (X:Orders.OrderedType) <: MSetInterface.RawSets X.
Include Ops X.
(** Generic definition of binary-search-trees and proofs of specifications for generic functions such as mem or fold. *)
Include MSetGenTree.Props X Color.
Local Notation Rd := (Node Red).
Local Notation Bk := (Node Black).
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.
(** ** Singleton set *)

Lemma singleton_spec x y : InT y (singleton x) <-> X.eq y x.
#[global]
Instance singleton_ok x : Ok (singleton x).