perf: replaceS and instantiateRevBetaS

This PR minimizes the number of expression allocations performed by
`replaceS` and `instantiateRevBetaS`.

src/Lean/Meta/Sym/AlphaShareBuilder.lean

@@ -144,4 +144,37 @@ def mkHaveS (x : Name) (t : Expr) (v : Expr) (b : Expr) : m Expr := do
     assertShared b
   share1 <| .letE x t v b true
 
+@[inline] def _root_.Lean.Expr.updateAppS! (e : Expr) (newFn : Expr) (newArg : Expr) : m Expr := do
+  let .app fn arg := e | panic! "application expected"
+  if isSameExpr fn newFn && isSameExpr arg newArg then return e else mkAppS newFn newArg
+
+@[inline] def _root_.Lean.Expr.updateMDataS! (e : Expr) (newExpr : Expr) : m Expr := do
+  let .mdata d a := e | panic! "mdata expected"
+  if isSameExpr a newExpr then return e else mkMDataS d newExpr
+
+@[inline] def _root_.Lean.Expr.updateProjS! (e : Expr) (newExpr : Expr) : m Expr := do
+  let .proj s i a := e | panic! "proj expected"
+  if isSameExpr a newExpr then return e else mkProjS s i newExpr
+
+@[inline] def _root_.Lean.Expr.updateForallS! (e : Expr) (newDomain : Expr) (newBody : Expr) : m Expr := do
+  let .forallE n d b bi := e | panic! "forall expected"
+  if isSameExpr d newDomain && isSameExpr b newBody then
+    return e
+  else
+    mkForallS n bi newDomain newBody
+
+@[inline] def _root_.Lean.Expr.updateLambdaS! (e : Expr) (newDomain : Expr) (newBody : Expr) : m Expr := do
+  let .lam n d b bi := e | panic! "lambda expected"
+  if isSameExpr d newDomain && isSameExpr b newBody then
+    return e
+  else
+    mkLambdaS n bi newDomain newBody
+
+@[inline] def _root_.Lean.Expr.updateLetS! (e : Expr) (newType : Expr) (newVal : Expr) (newBody : Expr) : m Expr := do
+  let .letE n t v b nondep := e | panic! "let expression expected"
+  if isSameExpr t newType && isSameExpr v newVal && isSameExpr b newBody then
+    return e
+  else
+    mkLetS n newType newVal newBody nondep
+
 end Lean.Meta.Sym.Internal

src/Lean/Meta/Sym/InstantiateS.lean

@@ -165,34 +165,39 @@ where
       if let some r := (← get)[key]? then
         return r
       else
-        save key (← mkAppS (← visitAppDefault f offset) (← visitChild a offset))
+        save key (← e.updateAppS! (← visitAppDefault f offset) (← visitChild a offset))
     | e => visitChild e offset
 
-  visitAppBeta (f : Expr) (argsRev : Array Expr) (offset : Nat) : M Expr := do
+  visitAppBeta (e : Expr) (f : Expr) (argsRev : Array Expr) (offset : Nat) (modified : Bool) : M Expr := do
     match f with
-    | .app f a => visitAppBeta f (argsRev.push (← visitChild a offset)) offset
+    | .app f a =>
+      let a' ← visitChild a offset
+      visitAppBeta e f (argsRev.push a') offset (modified || !isSameExpr a a')
     | .bvar bidx =>
-      let f ← visitBVar f bidx offset
-      betaRevS f argsRev
+      let f' ← visitBVar f bidx offset
+      if modified || !isSameExpr f f' then
+        betaRevS f' argsRev
+      else
+        return e
     | _ => unreachable!
 
-  visitApp (f : Expr) (arg : Expr) (offset : Nat) : M Expr := do
-    let arg ← visitChild arg offset
+  visitApp (e : Expr) (f : Expr) (arg : Expr) (offset : Nat) (_ : e = .app f arg) : M Expr := do
+    let arg' ← visitChild arg offset
     if f.getAppFn.isBVar then
-      visitAppBeta f #[arg] offset
+      visitAppBeta e f #[arg'] offset (!isSameExpr arg arg')
     else
-      mkAppS (← visitAppDefault f offset) arg
+      e.updateAppS! (← visitAppDefault f offset) arg'
 
   visit (e : Expr) (offset : Nat) : M Expr := do
-    match e with
+    match h : e with
     | .lit _ | .mvar _ | .fvar _ | .const _ _ | .sort _ => unreachable!
     | .bvar bidx => visitBVar e bidx offset
-    | .app f a => visitApp f a offset
-    | .mdata m a => mkMDataS m (← visitChild a offset)
-    | .proj s i a => mkProjS s i (← visitChild a offset)
-    | .forallE n d b bi => mkForallS n bi (← visitChild d offset) (← visitChild b (offset+1))
-    | .lam n d b bi => mkLambdaS n bi (← visitChild d offset) (← visitChild b (offset+1))
-    | .letE n t v b d => mkLetS n (← visitChild t offset) (← visitChild v offset) (← visitChild b (offset+1)) (nondep := d)
+    | .app f a => visitApp e f a offset h
+    | .mdata _ a => e.updateMDataS! (← visitChild a offset)
+    | .proj _ _ a => e.updateProjS! (← visitChild a offset)
+    | .forallE _ d b _ => e.updateForallS! (← visitChild d offset) (← visitChild b (offset+1))
+    | .lam _ d b _ => e.updateLambdaS! (← visitChild d offset) (← visitChild b (offset+1))
+    | .letE _ t v b _ => e.updateLetS! (← visitChild t offset) (← visitChild v offset) (← visitChild b (offset+1))
 
 /--
 Similar to `instantiateRevS`, but beta-reduces nested applications whose function becomes a lambda

src/Lean/Meta/Sym/ReplaceS.lean

@@ -33,12 +33,12 @@ mutual
 @[specialize] def visit (e : Expr) (offset : Nat) (fn : Expr → Nat → AlphaShareBuilderM (Option Expr)) : M Expr := do
   match e with
   | .lit _ | .mvar _ | .bvar _ | .fvar _ | .const _ _ | .sort _ => unreachable!
-  | .app f a => mkAppS (← visitChild f offset fn) (← visitChild a offset fn)
-  | .mdata m a => mkMDataS m (← visitChild a offset fn)
-  | .proj s i a => mkProjS s i (← visitChild a offset fn)
-  | .forallE n d b bi => mkForallS n bi (← visitChild d offset fn) (← visitChild b (offset+1) fn)
-  | .lam n d b bi => mkLambdaS n bi (← visitChild d offset fn) (← visitChild b (offset+1) fn)
-  | .letE n t v b d => mkLetS n (← visitChild t offset fn) (← visitChild v offset fn) (← visitChild b (offset+1) fn) (nondep := d)
+  | .app f a => e.updateAppS! (← visitChild f offset fn) (← visitChild a offset fn)
+  | .mdata _ a => e.updateMDataS! (← visitChild a offset fn)
+  | .proj _ _ a => e.updateProjS! (← visitChild a offset fn)
+  | .forallE _ d b _ => e.updateForallS! (← visitChild d offset fn) (← visitChild b (offset+1) fn)
+  | .lam _ d b _ => e.updateLambdaS! (← visitChild d offset fn) (← visitChild b (offset+1) fn)
+  | .letE _ t v b _ => e.updateLetS! (← visitChild t offset fn) (← visitChild v offset fn) (← visitChild b (offset+1) fn)
 end
 
 /--