From 38a1675c72b67bb455e954c0d00f1d23cb6d3a79 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Thu, 17 Mar 2022 16:44:38 -0700
Subject: [PATCH] fix: dicrimination tree `getMatchWithExtra`

It was buggy when the input argument could be reduced `e`.

fixes #1048
---
 src/Lean/Meta/DiscrTree.lean           | 27 ++++++++------------------
 src/Lean/Meta/Tactic/Simp/Rewrite.lean |  2 +-
 2 files changed, 9 insertions(+), 20 deletions(-)

diff --git a/src/Lean/Meta/DiscrTree.lean b/src/Lean/Meta/DiscrTree.lean
index 2b8c530694..8c8ac63663 100644
--- a/src/Lean/Meta/DiscrTree.lean
+++ b/src/Lean/Meta/DiscrTree.lean
@@ -507,26 +507,15 @@ partial def getMatch (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
   Similar to `getMatch`, but returns solutions that are prefixes of `e`.
   We store the number of ignored arguments in the result.-/
 partial def getMatchWithExtra (d : DiscrTree α) (e : Expr) : MetaM (Array (α × Nat)) :=
-  withReducible do
-    let result := getStarResult d |>.map (., 0)
-    let (k, args) ← getMatchKeyArgs e (root := true)
-    match k with
-    | Key.star => return result
-    | _        => process k args.toSubarray 0 result
+  go e 0 #[]
 where
-  process (k : Key) (args : Subarray Expr) (numExtraArgs : Nat) (result : Array (α × Nat)) : MetaM (Array (α × Nat)) := do
-    -- Remark: the args are stored in reverse order
-    let result :=
-      if d.root.find? k |>.isSome then
-        result ++ ((← getMatchRoot d k args.toArray #[]).map (., numExtraArgs))
-      else
-        result
-    match k with
-    | Key.const f 0     => return result
-    | Key.const f (n+1) => process (Key.const f n) args.popFront (numExtraArgs + 1) result
-    | Key.fvar f 0      => return result
-    | Key.fvar f (n+1)  => process (Key.fvar f n) args.popFront (numExtraArgs + 1) result
-    | _                 => return result
+  /- TODO: better implementation, this is one is very naive. -/
+  go (e : Expr) (numExtra : Nat) (result : Array (α × Nat)) : MetaM (Array (α × Nat)) := do
+    let result := result ++ (← getMatch d e).map (., numExtra)
+    if e.isApp then
+      go e.appFn! (numExtra + 1) result
+    else
+      return result
 
 partial def getUnify (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
   withReducible do
diff --git a/src/Lean/Meta/Tactic/Simp/Rewrite.lean b/src/Lean/Meta/Tactic/Simp/Rewrite.lean
index 7a6fa1050d..b539f9f310 100644
--- a/src/Lean/Meta/Tactic/Simp/Rewrite.lean
+++ b/src/Lean/Meta/Tactic/Simp/Rewrite.lean
@@ -63,7 +63,7 @@ private def tryTheoremCore (lhs : Expr) (xs : Array Expr) (bis : Array BinderInf
       if ← hasAssignableMVar proof then
         trace[Meta.Tactic.simp.rewrite] "{thm}, has unassigned metavariables after unification"
         return none
-      let rhs   ← instantiateMVars type.appArg!
+      let rhs := (← instantiateMVars type).appArg!
       if e == rhs then
         return none
       if thm.perm then