fix: simp was not applying rewrites to the function application prefixes

src/Lean/Meta/DiscrTree.lean

@@ -424,17 +424,32 @@ private def getStarResult (d : DiscrTree α) : Array α :=
 private abbrev findKey (cs : Array (Key × Trie α)) (k : Key) : Option (Key × Trie α) :=
   cs.binSearch (k, arbitrary) (fun a b => a.1 < b.1)
 
-partial def getMatch (d : DiscrTree α) (e : Expr) : MetaM (Array α) :=
+/--
+  Find values that match `e` in `d`.
+  If `allowExtraArgs == true`, we also return solutions that match prefixes of `e`.
+-/
+partial def getMatch (d : DiscrTree α) (e : Expr) (allowExtraArgs := false) : MetaM (Array α) :=
   withReducible do
     let result := getStarResult d
     let (k, args) ← getMatchKeyArgs e (root := true)
     match k with
     | Key.star => return result
-    | _        =>
-      match d.root.find? k with
-      | none   => return result
-      | some c => process args c result
+    | _        => if allowExtraArgs then processRootWithExtra k args result else processRoot k args result
 where
+  processRoot (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) := do
+    match d.root.find? k with
+    | none   => return result
+    | some c => process args c result
+
+  processRootWithExtra (k : Key) (args : Array Expr) (result : Array α) : MetaM (Array α) := do
+    let result ← processRoot k args result
+    match k with
+    | Key.const f 0     => return result
+    | Key.const f (n+1) => processRootWithExtra (Key.const f n) args.pop result
+    | Key.fvar f 0      => return result
+    | Key.fvar f (n+1)  => processRootWithExtra (Key.fvar f n) args.pop result
+    | _                 => return result
+
   process (todo : Array Expr) (c : Trie α) (result : Array α) : MetaM (Array α) := do
     match c with
     | Trie.node vs cs =>

src/Lean/Meta/Tactic/Simp/Rewrite.lean

@@ -3,6 +3,7 @@ Copyright (c) 2020 Microsoft Corporation. All rights reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Authors: Leonardo de Moura
 -/
+import Lean.Meta.AppBuilder
 import Lean.Meta.SynthInstance
 import Lean.Meta.Tactic.Simp.Types
 
@@ -48,34 +49,57 @@ def tryLemma? (e : Expr) (lemma : SimpLemma) (discharge? : Expr → SimpM (Optio
     let (xs, bis, type) ← forallMetaTelescopeReducing type
     let type ← whnf (← instantiateMVars type)
     let lhs := type.appFn!.appArg!
-    if (← isDefEq lhs e) then
-      unless (← synthesizeArgs lemma.getName xs bis discharge?) do
+    let rec go (e : Expr) : SimpM (Option Result) := do
+      if (← isDefEq lhs e) then
+        unless (← synthesizeArgs lemma.getName xs bis discharge?) do
+          return none
+        let proof ← instantiateMVars (mkAppN val xs)
+        if ← hasAssignableMVar proof then
+          trace[Meta.Tactic.simp.rewrite] "{lemma}, has unassigned metavariables after unification"
+          return none
+        let rhs   ← instantiateMVars type.appArg!
+        if e == rhs then
+          return none
+        if lemma.perm && !Expr.lt rhs e then
+          trace[Meta.Tactic.simp.rewrite] "{lemma}, perm rejected {e} ==> {rhs}"
+          return none
+        trace[Meta.Tactic.simp.rewrite] "{lemma}, {e} ==> {rhs}"
+        return some { expr := rhs, proof? := proof }
+      else
+        unless lhs.isMVar do
+          -- We do not report unification failures when `lhs` is a metavariable
+          -- Example: `x = ()`
+          -- TODO: reconsider if we want lemmas such as `(x : Unit) → x = ()`
+          trace[Meta.Tactic.simp.unify] "{lemma}, failed to unify {lhs} with {e}"
         return none
-      let proof ← instantiateMVars (mkAppN val xs)
-      if ← hasAssignableMVar proof then
-        trace[Meta.Tactic.simp.rewrite] "{lemma}, has unassigned metavariables after unification"
-        return none
-      let rhs   ← instantiateMVars type.appArg!
-      if e == rhs then
-        return none
-      if lemma.perm && !Expr.lt rhs e then
-        trace[Meta.Tactic.simp.rewrite] "{lemma}, perm rejected {e} ==> {rhs}"
-        return none
-      trace[Meta.Tactic.simp.rewrite] "{lemma}, {e} ==> {rhs}"
-      return some { expr := rhs, proof? := proof }
-    else
-      unless lhs.isMVar do
-        -- We do not report unification failures when `lhs` is a metavariable
-        -- Example: `x = ()`
-        -- TODO: reconsider if we want lemmas such as `(x : Unit) → x = ()`
-        trace[Meta.Tactic.simp.unify] "{lemma}, failed to unify {lhs} with {e}"
+    let lhsNumArgs := lhs.getAppNumArgs
+    let eNumArgs := e.getAppNumArgs
+    if eNumArgs == lhsNumArgs then
+      go e
+    else if eNumArgs < lhsNumArgs then
       return none
+    else
+      /- Check whether we need something more sophisticated here.
+         This simple approach was good enough for Mathlib 3 -/
+      let mut extraArgs := #[]
+      let mut e := e
+      for i in [:eNumArgs - lhsNumArgs] do
+        extraArgs := extraArgs.push e.appArg!
+        e := e.appFn!
+      match (← go e) with
+      | none => return none
+      | some { expr := eNew, proof? := none } => return some { expr := mkAppN eNew extraArgs }
+      | some { expr := eNew, proof? := some proof } =>
+        let mut proof := proof
+        for extraArg in extraArgs do
+          proof ← mkCongrFun proof extraArg
+        return some { expr := mkAppN eNew extraArgs, proof? := some proof }
 
 /-
 Remark: the parameter tag is used for creating trace messages. It is irrelevant otherwise.
 -/
 def rewrite (e : Expr) (s : DiscrTree SimpLemma) (erased : Std.PHashSet Name) (discharge? : Expr → SimpM (Option Expr)) (tag : String) : SimpM Result := do
-  let lemmas ← s.getMatch e
+  let lemmas ← s.getMatch e (allowExtraArgs := true)
   if lemmas.isEmpty then
     trace[Debug.Meta.Tactic.simp] "no theorems found for {tag}-rewriting {e}"
     return { expr := e }

tests/lean/simpPrefixIssue.lean

@@ -0,0 +1,8 @@
+universe u
+axiom f {α : Sort u} (a : α) : α
+axiom f_eq {α : Sort u} (a : α) : f a = a
+
+example (a : Nat) : f id a = a := by
+  simp only [f_eq]
+  traceState
+  rfl

tests/lean/simpPrefixIssue.lean.expected.out

@@ -0,0 +1,2 @@
+a : Nat
+⊢ id a = a