Module E := K.
Inductive color := Red | Black.
Inductive tree  : Type :=
| EE : tree
| T: color -> tree -> key -> tree -> tree.
Fixpoint member (x: key) (t : tree) : bool :=
match t with
| EE => false
| T _ tl k tr => match K.compare x k with
| Lt => member x tl
| Eq => true
| Gt => member x tr
end
end.
Definition balance color t1 k t2 :=
match color with Red => T Red t1 k t2
| _ =>
match t1, t2 with
| T Red (T Red a x b) y c, d => T Red (T Black a x b) y (T Black c k d)
| T Red a x (T Red b y c), d => T Red (T Black a x b) y (T Black c k d)
| a, T Red (T Red b y c) z d => T Red (T Black a k b) y (T Black c z d)
| a, T Red b y (T Red c z d) => T Red (T Black a k b) y (T Black c z d)
| _, _ => T Black t1 k t2
end
end.
Fixpoint ins x s :=
match s with
| EE => T Red EE x EE
| T c a y b => match K.compare x y with
| Lt => balance c (ins x a) y b
| Eq => T c a x b
| Gt => balance c a y (ins x b)
end
end.
Definition makeBlack t :=
match t with
| EE => EE
| T _ a x b => T Black a x b
end.
Definition insert x s := makeBlack (ins x s).
Fixpoint ins' x s :=
match s with
| EE => Some (T Red EE x EE)
| T c a y b => match K.compare x y with
| Lt => match ins' x a with
| None => None
| Some s' => Some (balance c s' y b)
end
| Eq => None
| Gt => match ins' x b with
| None => None
| Some s' => Some (balance c a y s')
end
end
end.
Definition mem_add' x s :=
match ins' x s with
| None => None
| Some s' => Some (makeBlack s')
end.
(* The next several definitions are to handle deletion, which is quite complex! *)
Definition balance' tl k tr :=
match tl , tr with
| T Red (T Red a x b) y c, d => T Red (T Black a x b) y (T Black c k d)
| T Red a x (T Red b y c), d => T Red (T Black a x b) y (T Black c k d)
| a, T Red b y (T Red c z d) => T Red (T Black a k b) y (T Black c z d)
| a, T Red (T Red b y c) z d => T Red (T Black a k b) y (T Black c z d)
| _ , _ => T Black tl k tr
end.
Definition sub1 t :=
match t with
| EE => EE
| T _ a x b => T Red a x b
end.
Definition balleft tl k tr :=
match tl with
| T Red a x b => T Red (T Black a x b) k tr
| _ => match tr with
| T Black a y b => balance' tl k (T Red a y b)
| T Red (T Black a y b) z c => T Red (T Black tl k a) y (balance' b z (sub1 c))
| _ => T Red tl k tr (* impossible *)
end
end.
Definition balright tl k tr :=
match tr with
| T Red b y c => T Red tl k (T Black b y c)
| _ => match tl with
| T Black a x b => balance' (T Red a x b) k tr
| T Red a x (T Black b y c) => T Red (balance' (sub1 a) x b) y (T Black c k tr)
| _ => T Red tl k tr (* impossible *)
end
end.
Fixpoint size t := match t with EE => 0 | T _ tl _ tr => 1 + size tl + size tr end.
Require Recdef Omega.
Function append (tlr: tree*tree) {measure (fun tlr => size (fst tlr) + size (snd tlr)) tlr} : tree :=
match tlr with
| (EE, x) => x
| (x, EE) => x
| (T Red a x b, T Red c y d) =>
match append (b,c) with
| T Red b' z c' => T Red (T Red a x b') z (T Red c' y d)
| bc => T Red a x (T Red bc y d)
end
| (T Black a x b, T Black c y d) =>
match append (b,c) with
| T Red b' z c' => T Red (T Black a x b') z (T Black c' y d)
| bc => balleft a x (T Black bc y d)
end
| (a, T Red b x c) => T Red (append(a,b)) x c
| (T Red a x b, c) => T Red a x (append(b,c))
end.
Definition ltopt (x y : option key) :=
match x, y with
| Some x', Some y' => K.lt x' y'
| _, _ => True
end.
Inductive searchtree: option key -> option key -> tree -> Prop :=
| STE: forall lo hi, ltopt lo hi -> searchtree lo hi EE
| STT: forall c tl k tr lo hi,
searchtree lo (Some k)  tl ->
searchtree (Some k) hi tr ->
searchtree lo hi (T c tl k tr) .
Ltac do_searchtree :=
assumption ||
constructor ||
match goal with
|  |- searchtree _ _ (match ?C with Red => _ | Black => _ end) => destruct C
|  |- searchtree _ _ (match ?C with EE => _ | T _ _ _ _ => _ end) => destruct C
| H: searchtree _ _ EE |- _ => inv H
| H: searchtree _ _ (T _ _ _ _) |- _ => inv H
| |- ltopt _ _ => unfold ltopt in *; auto
| |- match ?A with Some _ => _ | None => _ end => destruct A
| H: K.lt ?A ?B  |- K.lt ?A ?C  =>
try solve [apply lt_trans with B; assumption]; clear H
end.
Definition eqopt (a b : option key) :=
match a, b with
| Some x, Some y => K.eq x y
| _ , _ => False
end.
Inductive interp: tree -> (key -> Prop) :=
| member_here: forall x y c tl tr, K.eq x y -> interp (T c tl y tr) x
| member_left: forall x y c tl tr, interp tl x -> interp (T c tl y tr) x
| member_right: forall x y c tl tr, interp tr x -> interp (T c tl y tr) x.
Hint Constructors interp.
Ltac etac :=
match goal with
| H: interp EE _ |- _ => inv H
| H: InA _ _ nil |- _ => inv H
| H: false = true |- _ => inv H
| H: true = false |- _ => inv H
| _ => auto
end.
Ltac do_interp_balance :=
match goal with
| |- interp (match ?A with Red => _ | Black => _ end) _  => destruct A
| |- interp (match ?A with EE => _ | T _ _ _ _ => _ end) _ => destruct A
|  H: interp (T _ _ _ _) _ |- _ => inv H
| H: interp (match ?A with Red => _ | Black => _ end) _ |- _ => destruct A
| H: interp (match ?A with EE => _ | T _ _ _ _ => _ end) _ |- _ => destruct A
end; auto.
Inductive is_redblack : tree -> color -> nat -> Prop :=
| IsRB_leaf: forall c, is_redblack EE c 0
| IsRB_r: forall tl k tr n,
is_redblack tl Red n ->
is_redblack tr Red n ->
is_redblack (T Red tl k tr) Black n
| IsRB_b: forall c tl k tr n,
is_redblack tl Black n ->
is_redblack tr Black n ->
is_redblack (T Black tl k tr) c (S n).
Hint Resolve is_redblack_toblack.
Inductive infrared : tree -> nat -> Prop :=
(* This predicate expresses, "the tree is a red-black tree, except that it is permitted to have two red nodes in a row at the very root (only)."  *)
| infrared_e: infrared EE 0
| infrared_r: forall tl k tr n,
is_redblack tl Black n ->
is_redblack tr Black n ->
infrared (T Red tl k tr) n
| infrared_b: forall tl k tr n,
is_redblack tl Black n ->
is_redblack tr Black n ->
infrared (T Black tl k tr) (S n).
Ltac infrared_tac :=
repeat (
auto ||
constructor ||
match goal with
| |- infrared (match ?A with
| EE => _
| T _ _ _ _ => _
end)  _ =>
match goal with
| H: is_redblack A _ _ |- _ => inv H
| H: infrared A _ |- _ => inv H
end
| |- is_redblack (match ?A with
| EE => _
| T _ _ _ _ => _
end)  _ _ =>
match goal with
| H: is_redblack A _ _ |- _ => inv H
| H: infrared A _ |- _ => inv H
end
| |- infrared (match K.compare ?A ?B with Eq => _ | Lt => _ | Gt => _ end)  _ =>
do_compare A B
| |- is_redblack (match K.compare ?A ?B with Eq => _ | Lt => _ | Gt => _ end)  _ _ =>
do_compare A B
end).
(* This one is tedious with proof automation, but extremely tedious without proof automation *)
Definition valid (x: tree) := searchtree None None x /\
(exists n, is_redblack x Red n).
(*  NOW FOR THE HARD PART, DELETION! Here we follow the deletion algorithm of Stefan Kahrs (2001), and in particular, http://www.cs.kent.ac.uk/people/staff/smk/redblack/Untyped.hs *)
Ltac hax :=
repeat (auto; try solve [repeat constructor; auto];
match goal with
| H: is_redblack _ Red _ |- _ => inv H
| H: is_redblack _ _ 0 |- _ => inv H
| H: is_redblack _ _ (S _) |- _ => inv H
| H: infrared _ 0 |- _ => inv H
| H: infrared _ (S _) |- _ => inv H
end).

Lemma balance'_shape:
forall n tl k tr,
(is_redblack tl Black n /\ infrared tr n \/ infrared tl n /\ is_redblack tr Black n) ->
is_redblack (balance' tl k tr) Black (S n).