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).
Definition is_red_or_empty t :=
match t with T Black _ _ _ => False | _ => True end.
Definition is_black t :=
match t with T Black _ _ _ => True | _ => False end.
Hint Constructors interp.
Tactic Notation "do_hyp" hyp(H) :=
match type of H with ?a -> _ =>
let H1 := fresh in (assert (H1: a); [ |generalize (H H1); clear H H1; intro H]) end.
Ltac jax :=
repeat
match goal with
| H: interp EE _ |- _ => inv H
| |- interp match ?c with Red => _ | Black => _ end _ => destruct c
| |- interp match ?t with EE => _ | T _ _ _ _ => _ end _ =>
match t with balance' _ _ _ => fail 1 | append _ => fail 1 | _ =>  destruct t end
| H: interp (balance' _ _ _) _ |- _ =>
apply interp_balance' in H; destruct H as [?|[?|?]]
| |- interp (balance' _ _ _) _ =>
apply interp_balance'
| H: interp match ?c with Red => _ | Black => _ end _ |- _ => destruct c
| H: interp match ?t with EE => _ | T _ _ _ _ => _ end _ |- _  =>
match t with balance' _ _ _ => fail 1 | append _ => fail 1 | _ =>  destruct t end
| H: interp (T _ _ _ _) _ |- _ => inv H
| H: interp _ _ <-> interp _ _ \/ interp _ _ |- _ =>
solve [destruct H  as [H _]; destruct H; auto]
| |- _ => solve [auto 50]
| H: interp _ _ -> _ \/ _ |- _ \/ _ => solve [destruct H; auto]
| H: interp ?JJ  _, IHsz: forall _ _, _ -> _ |- _ =>
match JJ with context [append(?a,?b)] =>
let H99 := fresh in
destruct (IHsz a b) as [H99 _];  [simpl in *; omega | ];
destruct (append(a,b));
try  (destruct (H99 H); [ solve [auto] | | ])
end
|  IHsz: forall _ _, _ -> _ |- interp ?JJ _ =>
match JJ with context [append(?a,?b)] =>
let H99 := fresh in
destruct (IHsz a b) as [HA HB];  [simpl in *; omega | ];
destruct (append(a,b));
try (do_hyp HA; [solve [auto] | ]);
try (do_hyp HB; [solve [auto] | ])
end
end.
(*  To compare two trees for equality (or less/greater) is the classic "samefringe" problem.  I use a solution involving the "gopher function",  attributed by Gabriel 1991 to McCarthy 1977, but an equivalent solution was published first by Gruessay in 1976.  Gabriel writes, "this solution is difficult to understand," but in fact the Coq proofs about it are quite straightforward, actually simpler than the ones involving list-flattening solutions to samefringe.  Richard Gabriel, The design of parallel programming languages, in Artificial intelligence and mathematical theory of computation, ed. Vladimir Lifschitz,  pp. 91-108, Academic Press 1991.  Patrick Gruessay, "An iterative Lisp solution to the Samefringe problem", ACM SIGART Bulletin, Issue 59, p. 14, August 1976.  John McCarthy. Another Samefringe. ACM SIGART Bulletin, No. 61, February 1977.  *)
Fixpoint gopher (s: tree) : tree :=
match s with
| EE => EE
| T _ sl k sr => match gopher sl with
| T _ sll j slr => T Red sll j (T Red slr k sr)
| EE => s
end
end.

Lemma shape_gopher: forall a,
match gopher a with
| T _ (T _ _ _ _) _ _ => False
| _ => True
end.