From 35018dbea24dd92d5f75b02a9faa8bbf57a5496d Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sun, 10 Jul 2022 06:47:23 -0700
Subject: [PATCH] feat: unexpanders for `a[i]`, `a[i]' h`, `a[i]!`, and `a[i]?`

---
 src/Init/NotationExtra.lean          | 16 ++++++++++++++++
 src/Init/Tactics.lean                |  6 +++++-
 tests/lean/getElem.lean              |  9 +++++++++
 tests/lean/getElem.lean.expected.out | 10 ++++++++++
 4 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/src/Init/NotationExtra.lean b/src/Init/NotationExtra.lean
index 979cee0050a..2cf1faebf6b 100644
--- a/src/Init/NotationExtra.lean
+++ b/src/Init/NotationExtra.lean
@@ -160,6 +160,22 @@ macro "calc " steps:withPosition(calcStep+) : tactic => `(exact calc $steps*)
   | `($f [$k] $sep)  => `($f $k $sep)
   | _                => throw ()
 
+@[appUnexpander GetElem.getElem] def unexpandGetElem : Lean.PrettyPrinter.Unexpander
+  | `(getElem $array $index $_) => `($array[$index])
+  | _ => throw ()
+
+@[appUnexpander getElem!] def unexpandGetElem! : Lean.PrettyPrinter.Unexpander
+  | `(getElem! $array $index) => `($array[$index]!)
+  | _ => throw ()
+
+@[appUnexpander getElem?] def unexpandGetElem? : Lean.PrettyPrinter.Unexpander
+  | `(getElem? $array $index) => `($array[$index]?)
+  | _ => throw ()
+
+@[appUnexpander getElem'] def unexpandGetElem' : Lean.PrettyPrinter.Unexpander
+  | `(getElem' $array $index $h) => `($array[$index]'$h)
+  | _ => throw ()
+
 /--
 Apply function extensionality and introduce new hypotheses.
 The tactic `funext` will keep applying new the `funext` lemma until the goal target is not reducible to
diff --git a/src/Init/Tactics.lean b/src/Init/Tactics.lean
index f0296142cf9..7b822b3d677 100644
--- a/src/Init/Tactics.lean
+++ b/src/Init/Tactics.lean
@@ -442,4 +442,8 @@ macro "get_elem_tactic" : tactic =>
 
 macro:max x:term noWs "[" i:term "]" : term => `(getElem $x $i (by get_elem_tactic))
 
-macro x:term noWs "[" i:term "]'" h:term:max : term => `(getElem $x $i $h)
+/-- Helper declaration for the unexpander -/
+@[inline] def getElem' [GetElem Cont Idx Elem Dom] (xs : Cont) (i : Idx) (h : Dom xs i) : Elem :=
+  getElem xs i h
+
+macro x:term noWs "[" i:term "]'" h:term:max : term => `(getElem' $x $i $h)
diff --git a/tests/lean/getElem.lean b/tests/lean/getElem.lean
index 242ed2cbd4e..8dbb6d3f9b4 100644
--- a/tests/lean/getElem.lean
+++ b/tests/lean/getElem.lean
@@ -20,3 +20,12 @@ def f5 (i : Nat) (h : i < n) :=
 
 def f6 (i : Nat) :=
   a[i]!
+
+def f7 (i : Nat) :=
+  a[i]?
+
+#print f2
+#print f3
+#print f5
+#print f6
+#print f7
diff --git a/tests/lean/getElem.lean.expected.out b/tests/lean/getElem.lean.expected.out
index b731036abf2..bba4aeb838c 100644
--- a/tests/lean/getElem.lean.expected.out
+++ b/tests/lean/getElem.lean.expected.out
@@ -6,3 +6,13 @@ getElem.lean:2:2-2:6: error: failed to prove index is valid, possible solutions:
 a : Array Nat
 i : Nat
 ⊢ i < Array.size a
+def f2 : (a : Array Nat) → Fin (Array.size a) → Nat :=
+fun a i => a[i]
+def f3 : {n : Nat} → (a : Array Nat) → n ≤ Array.size a → Fin n → Nat :=
+fun {n} a h i => a[i]
+def f5 : (i : Nat) → i < n → Nat :=
+fun i h => a[i]'(_ : i < Array.size a)
+def f6 : Nat → Nat :=
+fun i => a[i]!
+def f7 : Nat → Option Nat :=
+fun i => a[i]?