feat(library/data/rbtree): add find and define contains using it

library/data/rbtree/find.lean

@@ -13,7 +13,7 @@ namespace rbnode
 variables {α : Type u}
 
 @[elab_simple]
-lemma contains.induction {p : rbnode α → Prop} (lt) [decidable_rel lt]
+lemma find.induction {p : rbnode α → Prop} (lt) [decidable_rel lt]
    (t x)
    (h₁ : p leaf)
    (h₂ : ∀ l y r (h : cmp_using lt x y = ordering.lt) (ih : p l), p (red_node l y r))
@@ -42,10 +42,10 @@ begin
   }
 end
 
-lemma contains_correct {t : rbnode α} {lt x} [decidable_rel lt] [is_strict_weak_order α lt] : ∀ {lo hi} (hs : is_searchable lt t lo hi), mem lt x t ↔ contains lt t x = tt :=
+lemma find_correct {t : rbnode α} {lt x} [decidable_rel lt] [is_strict_weak_order α lt] : ∀ {lo hi} (hs : is_searchable lt t lo hi), mem lt x t ↔ ∃ y, find lt t x = some y ∧ x ≈[lt] y :=
 begin
-  apply contains.induction lt t x; intros; simp only [mem, contains, *],
-  { simp },
+  apply find.induction lt t x; intros; simp only [mem, find, *],
+  { simp, intro h, cases h with _ h, cases h, contradiction },
   twice { -- red and black cases are identical
     {
       cases hs,
@@ -63,7 +63,7 @@ begin
       },
       { intro hc, left, exact iff.mpr (ih hs₁) hc },
     },
-    { simp at h, simp [h] },
+    { simp at h, simp [h, strict_weak_order.equiv], existsi y, split, refl, assumption },
     {
       cases hs,
       apply iff.intro,
@@ -82,4 +82,14 @@ begin
     } }
 end
 
+lemma eqv_of_find_some {t : rbnode α} {lt x y} [decidable_rel lt] [is_strict_weak_order α lt] : ∀ {lo hi} (hs : is_searchable lt t lo hi) (he : find lt t x = some y), x ≈[lt] y :=
+begin
+  apply find.induction lt t x; intros; simp only [mem, find, *] at *,
+  { contradiction },
+  twice {
+    { cases hs, exact ih hs₁ rfl },
+    { injection he, subst y, simp at h, exact h },
+    { cases hs, exact ih hs₂ rfl } }
+end
+
 end rbnode

library/data/rbtree/main.lean

@@ -3,7 +3,7 @@ Copyright (c) 2017 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
-import data.rbtree.contains data.rbtree.insert data.rbtree.min_max
+import data.rbtree.find data.rbtree.insert data.rbtree.min_max
 universes u
 
 /- TODO(Leo): remove after we cleanup stdlib simp lemmas -/
@@ -29,8 +29,21 @@ by simp [has_mem.mem, rbtree.mem, rbnode.mem, mk_rbtree]
 
 variables [decidable_rel lt]
 
+lemma find_correct [is_strict_weak_order α lt] (a : α) (t : rbtree α lt) : a ∈ t ↔ (∃ b, t.find a = some b ∧ a ≈[lt] b) :=
+begin cases t, apply rbnode.find_correct, apply rbnode.is_searchable_of_well_formed, assumption end
+
+lemma eqv_of_find_some [is_strict_weak_order α lt] {a b : α} {t : rbtree α lt} : t.find a = some b → a ≈[lt] b :=
+begin cases t, apply rbnode.eqv_of_find_some, apply rbnode.is_searchable_of_well_formed, assumption end
+
 lemma contains_correct [is_strict_weak_order α lt] (a : α) (t : rbtree α lt) : a ∈ t ↔ (t.contains a = tt) :=
-begin cases t, apply rbnode.contains_correct, apply rbnode.is_searchable_of_well_formed, assumption end
+begin
+  have h := find_correct a t,
+  simp [h, contains], apply iff.intro,
+  { intro h', cases h' with _ h', cases h', simp [*], simp [option.is_some] },
+  { intro h',
+    generalize heq : find t a = s, cases s with v, simp [heq, option.is_some] at h', contradiction,
+    existsi v, simp, apply eqv_of_find_some heq }
+end
 
 lemma mem_insert_of_incomp {a b : α} (t : rbtree α lt) : (¬ lt a b ∧ ¬ lt b a) → a ∈ t.insert b :=
 begin cases t, apply rbnode.mem_insert_of_incomp end

library/init/data/rbtree/basic.lean

@@ -122,19 +122,19 @@ def mem : α → rbnode α → Prop
 
 variable [decidable_rel lt]
 
-def contains : rbnode α → α → bool
-| leaf             x := ff
+def find : rbnode α → α → option α
+| leaf             x := none
 | (red_node a y b) x :=
   match cmp_using lt x y with
-  | ordering.lt := contains a x
-  | ordering.eq := tt
-  | ordering.gt := contains b x
+  | ordering.lt := find a x
+  | ordering.eq := some y
+  | ordering.gt := find b x
   end
 | (black_node a y b) x :=
   match cmp_using lt x y with
-  | ordering.lt := contains a x
-  | ordering.eq := tt
-  | ordering.gt := contains b x
+  | ordering.lt := find a x
+  | ordering.eq := some y
+  | ordering.gt := find b x
   end
 
 end membership
@@ -181,8 +181,11 @@ variables [decidable_rel lt]
 def insert : rbtree α lt → α → rbtree α lt
 | ⟨t, w⟩   x := ⟨t.insert lt x, well_formed.insert_wff w rfl⟩
 
-def contains : rbtree α lt → α → bool
-| ⟨t, _⟩ x := t.contains lt x
+def find : rbtree α lt → α → option α
+| ⟨t, _⟩ x := t.find lt x
+
+def contains (t : rbtree α lt) (a : α) : bool :=
+(t.find a).is_some
 
 def from_list (l : list α) (lt : α → α → Prop . rbtree.default_lt) [decidable_rel lt] : rbtree α lt :=
 l.foldl insert (mk_rbtree α lt)