From d551ff653d9046fe9970b238751ac259db1b2de1 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 15 Dec 2019 17:33:36 -0800
Subject: [PATCH] feat: add `getOp` for `a[i]` notation

---
 src/Init/Data/Array/Basic.lean             | 3 +++
 src/Init/Data/HashMap/Basic.lean           | 3 +++
 src/Init/Data/PersistentArray/Basic.lean   | 3 +++
 src/Init/Data/PersistentHashMap/Basic.lean | 3 +++
 4 files changed, 12 insertions(+)

diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean
index c5127377d5..1b91abfecc 100644
--- a/src/Init/Data/Array/Basic.lean
+++ b/src/Init/Data/Array/Basic.lean
@@ -93,6 +93,9 @@ if h : i < a.size then some (a.get ⟨i, h⟩) else none
 def getD (a : Array α) (i : Nat) (v₀ : α) : α :=
 if h : i < a.size then a.get ⟨i, h⟩ else v₀
 
+def getOp [Inhabited α] (self : Array α) (idx : Nat) : α :=
+self.get! idx
+
 @[extern "lean_array_fset"]
 def set (a : Array α) (i : @& Fin a.size) (v : α) : Array α :=
 { sz   := a.sz,
diff --git a/src/Init/Data/HashMap/Basic.lean b/src/Init/Data/HashMap/Basic.lean
index 67937ffd84..2cf5c26457 100644
--- a/src/Init/Data/HashMap/Basic.lean
+++ b/src/Init/Data/HashMap/Basic.lean
@@ -154,6 +154,9 @@ match m.find a with
 | some b => b
 | none   => panic! "key is not in the map"
 
+@[inline] def getOp (self : HashMap α β) (idx : α) : Option β :=
+self.find idx
+
 @[inline] def contains (m : HashMap α β) (a : α) : Bool :=
 match m with
 | ⟨ m, _ ⟩ => m.contains a
diff --git a/src/Init/Data/PersistentArray/Basic.lean b/src/Init/Data/PersistentArray/Basic.lean
index 7a71c2bf2b..a41573e834 100644
--- a/src/Init/Data/PersistentArray/Basic.lean
+++ b/src/Init/Data/PersistentArray/Basic.lean
@@ -68,6 +68,9 @@ if i >= t.tailOff then
 else
   getAux t.root (USize.ofNat i) t.shift
 
+def getOp [Inhabited α] (self : PersistentArray α) (idx : Nat) : α :=
+self.get! idx
+
 partial def setAux : PersistentArrayNode α → USize → USize → α → PersistentArrayNode α
 | node cs, i, shift, a =>
   let j     := div2Shift i shift;
diff --git a/src/Init/Data/PersistentHashMap/Basic.lean b/src/Init/Data/PersistentHashMap/Basic.lean
index b7a22fc817..317a58ed3e 100644
--- a/src/Init/Data/PersistentHashMap/Basic.lean
+++ b/src/Init/Data/PersistentHashMap/Basic.lean
@@ -155,6 +155,9 @@ partial def findAux [HasBeq α] : Node α β → USize → α → Option β
 def find [HasBeq α] [Hashable α] : PersistentHashMap α β → α → Option β
 | { root := n, .. }, k => findAux n (hash k) k
 
+@[inline] def getOp [HasBeq α] [Hashable α] (self : PersistentHashMap α β) (idx : α) : Option β :=
+self.find idx
+
 @[inline] def findD [HasBeq α] [Hashable α] (m : PersistentHashMap α β) (a : α) (b₀ : β) : β :=
 (m.find a).getD b₀