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.
Fixpoint rgopher (s: tree) : tree :=
match s with
| EE => EE
| T _ sl k sr => match rgopher sr with
| T _ srl j srr => T Red  (T Red sl k srl) j srr
| EE => s
end
end.
Function compare' (s t : tree) {measure size s} : comparison :=
match gopher s, gopher t with
| EE , EE => Eq
| T _ EE j sr , T _ EE j' tr =>
match K.compare j j' with
| Gt => Gt | Lt => Lt | Eq => compare' sr tr
end
| T _ _ _ _, T _ _ _ _ => Eq (* impossible *)
| T _ _ _ _, EE => Gt
|  EE, T _ _ _ _ => Lt
end.
Fixpoint revsize (n: positive) (lr l: list key) : positive * list key :=
match l with
| nil => (n, lr)
| x::xs => revsize (Psucc n) (x::lr) xs
end.
Function linear_union_aux (st : tree*tree) (nl: positive * list key)
{measure (fun st => size (fst st) + size (snd st)) st} : positive * list key  :=
match rgopher (fst st), rgopher (snd st) with
| EE , EE => nl
| T _ sl j EE as st1 , T _ tl j' EE as st2 =>
match K.compare j j' with
| Lt => linear_union_aux (st1, tl) (Psucc (fst nl), j':: snd nl)
| Gt => linear_union_aux (sl, st2) (Psucc (fst nl), j:: snd nl)
| Eq => linear_union_aux (sl,tl) (Psucc (fst nl), j:: snd nl)
end
| T _ sl j EE, EE => linear_union_aux (sl,EE) (Psucc (fst nl), j:: snd nl)
|  EE, T _ tl j' EE => linear_union_aux (EE,tl) (Psucc (fst nl), j':: snd nl)
| _, _ => nl (* impossible *)
end.
Ltac spec H := match type of H with ?A -> _ =>
let H' := fresh in (assert (H' : A); [ | specialize (H H'); clear H']) end.
Definition HdRel' (R: key -> key -> Prop) (lo: option key) (l: list key) :=
match lo with Some lo' => HdRel R lo' l | None => True end.
Fixpoint filter_aux (f: key -> bool) (s: tree) (n: positive) (l: list key) : (positive * list key) :=
match s with
| T _ sl k sr => let (n,l') := filter_aux f sr n l
in if f k then filter_aux f sl (Psucc n) (k::l')
else filter_aux f sl n l'
| EE => (n,l)
end.
Definition Kle x y := K.lt x y \/ K.eq x y.
induction 1; constructor. left; auto.
Qed.
Notation compatb := (Proper (E.eq==>Logic.eq)) (only parsing).
Definition filter' (f: key -> bool) (s: tree) : tree :=
let nl := filter_aux f s 1%positive nil in fst (treeify_g (fst nl) (snd nl)).
Fixpoint partition_aux (f: key -> bool) (s: tree)
(nl1 nl2: positive *  list key) :  (positive * list key) * (positive * list key) :=
match s with
| T _ sl k sr => match partition_aux f sr nl1 nl2
with (nl1', nl2') =>
if f k then partition_aux f sl (Psucc (fst nl1'), k:: snd nl1') nl2'
else partition_aux f sl nl1' (Psucc (fst nl2'), k:: snd nl2')
end
| EE => (nl1,nl2)
end.
Definition partition'  (f: key -> bool) (s: tree) : (tree*tree) :=
match partition_aux f s (1%positive, nil) (1%positive,  nil)
with ((n1,l1),(n2,l2)) => (fst (treeify_g n1 l1), fst (treeify_g n2 l2))
end.
induction s; intros; auto.
simpl fold'.
remember (fold' (list key) (@cons key)) as g.
rewrite <- Heqg.
rewrite <- IHs2.
f_equal. simpl.
f_equal.
apply IHs1.
Qed.
Definition union' (s1 s2: tree) : tree :=
match compare_height s1 s1 s2 s2 with
| Lt => fold' _ insert s1 s2
| Gt => fold' _ insert s2 s1
| Eq => linear_union s1 s2
end.
Definition diff' (s1 s2: tree) : tree :=
match compare_height s1 s1 s2 s2 with
| Lt => filter' (fun k => negb (member k s2)) s1
| Gt => fold' _ delete s2 s1
| Eq => linear_diff s1 s2
end.
Definition inter' (s1 s2: tree) : tree :=
match compare_height s1 s1 s2 s2 with
| Lt => filter' (fun k => member k s2) s1
| Gt => filter' (fun k => member k s1) s2
| Eq => linear_inter s1 s2
end.
Fixpoint elements' (x: tree) (base: list key) : list key :=
match x with
| EE => base
| T _ t1 k t2 => elements' t1 (cons k (elements' t2 base))
end.

Lemma filter'_spec : forall s x f,
compatb f ->
valid s ->
(interp (filter' f s) x <-> interp s x /\ f x = true).