(** The type of color annotation. *)

Inductive color := Red | Black.(** The type of color annotation. *)

Inductive color := Red | Black.(** The type of color annotation. *)

Inductive color := Red | Black.Definition t := 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.(** ** Red-Black invariants *)

(** In a red-black tree :
    - a red node has no red children
    - the black depth at each node is the same along all paths.
    The black depth is here an argument of the predicate. *)

Inductive rbt : nat -> tree -> Prop :=
 | RB_Leaf : rbt 0 Leaf
 | RB_Rd n l k r :
   notred l -> notred r -> rbt n l -> rbt n r -> rbt n (Rd l k r)
 | RB_Bk n l k r : rbt n l -> rbt n r -> rbt (S n) (Bk l k r).(** ** Red-Black invariants *)

(** In a red-black tree :
    - a red node has no red children
    - the black depth at each node is the same along all paths.
    The black depth is here an argument of the predicate. *)

Inductive rbt : nat -> tree -> Prop :=
 | RB_Leaf : rbt 0 Leaf
 | RB_Rd n l k r :
   notred l -> notred r -> rbt n l -> rbt n r -> rbt n (Rd l k r)
 | RB_Bk n l k r : rbt n l -> rbt n r -> rbt (S n) (Bk l k r).(** ** Red-Black invariants *)

(** In a red-black tree :
    - a red node has no red children
    - the black depth at each node is the same along all paths.
    The black depth is here an argument of the predicate. *)

Inductive rbt : nat -> tree -> Prop :=
 | RB_Leaf : rbt 0 Leaf
 | RB_Rd n l k r :
   notred l -> notred r -> rbt n l -> rbt n r -> rbt n (Rd l k r)
 | RB_Bk n l k r : rbt n l -> rbt n r -> rbt (S n) (Bk l k r).(** ** Red-Black invariants *)

(** In a red-black tree :
    - a red node has no red children
    - the black depth at each node is the same along all paths.
    The black depth is here an argument of the predicate. *)

Inductive rbt : nat -> tree -> Prop :=
 | RB_Leaf : rbt 0 Leaf
 | RB_Rd n l k r :
   notred l -> notred r -> rbt n l -> rbt n r -> rbt n (Rd l k r)
 | RB_Bk n l k r : rbt n l -> rbt n r -> rbt (S n) (Bk l k r).(** A red-red tree is almost a red-black tree, except that it has
    a _red_ root node which _may_ have red children. Note that a
    red-red tree is hence non-empty, and all its strict subtrees
    are red-black. *)

Inductive rrt (n:nat) : tree -> Prop :=
 | RR_Rd l k r : rbt n l -> rbt n r -> rrt n (Rd l k r).(** A red-red tree is almost a red-black tree, except that it has
    a _red_ root node which _may_ have red children. Note that a
    red-red tree is hence non-empty, and all its strict subtrees
    are red-black. *)

Inductive rrt (n:nat) : tree -> Prop :=
 | RR_Rd l k r : rbt n l -> rbt n r -> rrt n (Rd l k r).(** An almost-red-black tree is almost a red-black tree, except that
    it's permitted to have two red nodes in a row at the very root (only).
    We implement this notion by saying that a quasi-red-black tree
    is either a red-black tree or a red-red tree. *)

Inductive arbt (n:nat)(t:tree) : Prop :=
 | ARB_RB : rbt n t -> arbt n t
 | ARB_RR : rrt n t -> arbt n t.(** An almost-red-black tree is almost a red-black tree, except that
    it's permitted to have two red nodes in a row at the very root (only).
    We implement this notion by saying that a quasi-red-black tree
    is either a red-black tree or a red-red tree. *)

Inductive arbt (n:nat)(t:tree) : Prop :=
 | ARB_RB : rbt n t -> arbt n t
 | ARB_RR : rrt n t -> arbt n t.(** An almost-red-black tree is almost a red-black tree, except that
    it's permitted to have two red nodes in a row at the very root (only).
    We implement this notion by saying that a quasi-red-black tree
    is either a red-black tree or a red-red tree. *)

Inductive arbt (n:nat)(t:tree) : Prop :=
 | ARB_RB : rbt n t -> arbt n t
 | ARB_RR : rrt n t -> arbt n t.(** The main exported invariant : being a red-black tree for some
    black depth. *)

Class Rbt (t:tree) :=  RBT : exists d, rbt d t.(** The main exported invariant : being a red-black tree for some
    black depth. *)

Class Rbt (t:tree) :=  RBT : exists d, rbt d t.(** ** Basic tactics and results about red-black *)

Scheme rbt_ind := Induction for rbt Sort Prop.Ltac invrb := intros; invtree rrt; invtree rbt; try contradiction.Ltac desarb := match goal with H:arbt _ _ |- _ => destruct H end.Ltac nonzero n := destruct n as [|n]; [try split; invrb|].

Lemma arb_nr_rb n t : arbt n t -> notred t -> rbt n t.