From 338c9359fa4de2c8198256a6ffd8df463f009f0c Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 3 Oct 2019 09:32:53 -0700
Subject: [PATCH] feat: add `find!` to maps

---
 library/init/data/hashmap/basic.lean          | 11 +++++--
 .../init/data/persistenthashmap/basic.lean    |  5 ++++
 library/init/data/rbmap/basicaux.lean         | 30 +++++++++++++++++++
 library/init/data/rbmap/default.lean          |  3 +-
 library/init/lean/smap.lean                   |  5 ++++
 5 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 library/init/data/rbmap/basicaux.lean

diff --git a/library/init/data/hashmap/basic.lean b/library/init/data/hashmap/basic.lean
index f5a7868dff..cf551310ec 100644
--- a/library/init/data/hashmap/basic.lean
+++ b/library/init/data/hashmap/basic.lean
@@ -4,8 +4,10 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Author: Leonardo de Moura
 -/
 prelude
-import init.data.array.basic init.data.assoclist
-import init.data.option.basic init.data.hashable
+import init.data.array.basic
+import init.data.assoclist
+import init.data.option.basic
+import init.data.hashable
 universes u v w
 
 def HashMapBucket (α : Type u) (β : Type v) :=
@@ -147,6 +149,11 @@ match m with
 @[inline] def findD (m : HashMap α β) (a : α) (b₀ : β) : β :=
 (m.find a).getOrElse b₀
 
+@[inline] def find! [Inhabited β] (m : HashMap α β) (a : α) : β :=
+match m.find a with
+| some b => b
+| none   => panic! "key is not in the map"
+
 @[inline] def contains (m : HashMap α β) (a : α) : Bool :=
 match m with
 | ⟨ m, _ ⟩ => m.contains a
diff --git a/library/init/data/persistenthashmap/basic.lean b/library/init/data/persistenthashmap/basic.lean
index 46849b1940..cd2942ad5c 100644
--- a/library/init/data/persistenthashmap/basic.lean
+++ b/library/init/data/persistenthashmap/basic.lean
@@ -158,6 +158,11 @@ def find [HasBeq α] [Hashable α] : PersistentHashMap α β → α → Option 
 @[inline] def findD [HasBeq α] [Hashable α] (m : PersistentHashMap α β) (a : α) (b₀ : β) : β :=
 (m.find a).getOrElse b₀
 
+@[inline] def find! [HasBeq α] [Hashable α] [Inhabited β] (m : PersistentHashMap α β) (a : α) : β :=
+match m.find a with
+| some b => b
+| none   => panic! "key is not in the map"
+
 partial def containsAtAux [HasBeq α] (keys : Array α) (vals : Array β) (heq : keys.size = vals.size) : Nat → α → Bool
 | i, k =>
   if h : i < keys.size then
diff --git a/library/init/data/rbmap/basicaux.lean b/library/init/data/rbmap/basicaux.lean
new file mode 100644
index 0000000000..5a0695653d
--- /dev/null
+++ b/library/init/data/rbmap/basicaux.lean
@@ -0,0 +1,30 @@
+/-
+Copyright (c) 2019 Microsoft Corporation. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+prelude
+import init.data.rbmap.basic
+import init.util
+
+universes u v w w'
+
+namespace RBMap
+variables {α : Type u} {β : Type v} {lt : α → α → Bool}
+
+@[inline] def min! [Inhabited α] [Inhabited β] (t : RBMap α β lt) : α × β :=
+match t.min with
+| some p => p
+| none   => panic! "map is empty"
+
+@[inline] def max! [Inhabited α] [Inhabited β] (t : RBMap α β lt) : α × β :=
+match t.max with
+| some p => p
+| none   => panic! "map is empty"
+
+@[inline] def find! [Inhabited β] (t : RBMap α β lt) (k : α) : β :=
+match t.find k with
+| some b => b
+| none   => panic! "key is not in the map"
+
+end RBMap
diff --git a/library/init/data/rbmap/default.lean b/library/init/data/rbmap/default.lean
index 4d5a38b69a..0b9431cf8c 100644
--- a/library/init/data/rbmap/default.lean
+++ b/library/init/data/rbmap/default.lean
@@ -4,4 +4,5 @@ Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
 prelude
-import init.data.rbtree init.data.rbmap.basic
+import init.data.rbmap.basic
+import init.data.rbmap.basicaux
diff --git a/library/init/lean/smap.lean b/library/init/lean/smap.lean
index 9c68e51974..bf5533615a 100644
--- a/library/init/lean/smap.lean
+++ b/library/init/lean/smap.lean
@@ -49,6 +49,11 @@ instance : HasEmptyc (SMap α β) := ⟨SMap.empty⟩
 @[inline] def findD (m : SMap α β) (a : α) (b₀ : β) : β :=
 (m.find a).getOrElse b₀
 
+@[inline] def find! [Inhabited β] (m : SMap α β) (a : α) : β :=
+match m.find a with
+| some b => b
+| none   => panic! "key is not in the map"
+
 @[specialize] def contains : SMap α β → α → Bool
 | ⟨true, m₁, _⟩, k   => m₁.contains k
 | ⟨false, m₁, m₂⟩, k => m₁.contains k || m₂.contains k