From 41b491483610a030b67535c73ccfc718776a81af Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Wed, 17 Jul 2024 18:17:47 +0200
Subject: [PATCH] perf: `Replacement.apply` (#4776)

Avoid potentially expensive `e.replace` if it is not applicable.
---
 src/Lean/Data/RBMap.lean     |  7 +++++++
 src/Lean/Elab/MutualDef.lean | 24 +++++++++++++++++++-----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/src/Lean/Data/RBMap.lean b/src/Lean/Data/RBMap.lean
index 95cdd329dc..38200e5f19 100644
--- a/src/Lean/Data/RBMap.lean
+++ b/src/Lean/Data/RBMap.lean
@@ -80,6 +80,10 @@ protected def max : RBNode α β → Option (Sigma (fun k => β k))
 def singleton (k : α) (v : β k) : RBNode α β :=
   node red leaf k v leaf
 
+def isSingleton : RBNode α β → Bool
+  | node _ leaf _ _ leaf => true
+  | _ => false
+
 -- the first half of Okasaki's `balance`, concerning red-red sequences in the left child
 @[inline] def balance1 : RBNode α β → (a : α) → β a → RBNode α β → RBNode α β
   | node red (node red a kx vx b) ky vy c, kz, vz, d
@@ -269,6 +273,9 @@ variable {α : Type u} {β : Type v} {σ : Type w} {cmp : α → α → Ordering
 def depth (f : Nat → Nat → Nat) (t : RBMap α β cmp) : Nat :=
   t.val.depth f
 
+def isSingleton (t : RBMap α β cmp) : Bool :=
+  t.val.isSingleton
+
 @[inline] def fold (f : σ → α → β → σ) : (init : σ) → RBMap α β cmp → σ
   | b, ⟨t, _⟩ => t.fold f b
 
diff --git a/src/Lean/Elab/MutualDef.lean b/src/Lean/Elab/MutualDef.lean
index c4884048de..73aa6bf4ed 100644
--- a/src/Lean/Elab/MutualDef.lean
+++ b/src/Lean/Elab/MutualDef.lean
@@ -728,12 +728,26 @@ def insertReplacementForLetRecs (r : Replacement) (letRecClosures : List LetRecC
   letRecClosures.foldl (init := r) fun r c =>
     r.insert c.toLift.fvarId c.closed
 
+def isApplicable (r : Replacement) (e : Expr) : Bool :=
+  Option.isSome <| e.findExt? fun e =>
+    if e.hasFVar then
+      match e with
+      | .fvar fvarId => if r.contains fvarId then .found else .done
+      | _ => .visit
+    else
+      .done
+
 def Replacement.apply (r : Replacement) (e : Expr) : Expr :=
-  e.replace fun e => match e with
-    | .fvar fvarId => match r.find? fvarId with
-      | some c => some c
-      | _      => none
-    | _ => none
+  -- Remark: if `r` is not a singlenton, then declaration is using `mutual` or `let rec`,
+  -- and there is a big chance `isApplicable r e` is true.
+  if r.isSingleton && !isApplicable r e then
+    e
+  else
+    e.replace fun e => match e with
+      | .fvar fvarId => match r.find? fvarId with
+        | some c => some c
+        | _      => none
+      | _ => none
 
 def pushMain (preDefs : Array PreDefinition) (sectionVars : Array Expr) (mainHeaders : Array DefViewElabHeader) (mainVals : Array Expr)
     : TermElabM (Array PreDefinition) :=