perf: cache inside matcher WHNFs at reducible/instances/implicit

src/Lean/Meta/Basic.lean

@@ -460,6 +460,17 @@ register_builtin_option maxSynthPendingDepth : Nat := {
   descr    := "maximum number of nested `synthPending` invocations. When resolving unification constraints, pending type class problems may need to be synthesized. These type class problems may create new unification constraints that again require solving new type class problems. This option puts a threshold on how many nested problems are created."
 }
 
+inductive UnfoldPred where
+| standard : UnfoldPred
+| atMatcher : (Config → ConstantInfo → CoreM Bool) → UnfoldPred -- slightly dangerous; we ignore the function in the cache key, but it's always the same
+| customUncached : (Config → ConstantInfo → CoreM Bool) → UnfoldPred
+
+def UnfoldPred.isCacheable (p : UnfoldPred) : Bool :=
+  p matches .standard | .atMatcher _
+
+def UnfoldPred.cacheKey (p : UnfoldPred) : Bool :=
+  p matches .standard
+
 /--
   Contextual information for the `MetaM` monad.
 -/
@@ -503,7 +514,7 @@ structure Context where
   /--
     A predicate to control whether a constant can be unfolded or not at `whnf`.
     Note that we do not cache results at `whnf` when `canUnfold?` is not `none`. -/
-  canUnfold?        : Option (Config → ConstantInfo → CoreM Bool) := none
+  canUnfold?        : UnfoldPred := .standard
   /--
   When `Config.univApprox := true`, this flag is set to `true` when there is no
   progress processing universe constraints.
@@ -523,7 +534,7 @@ structure Context where
 deriving Inhabited
 
 def Context.config (c : Context) : Config := c.keyedConfig.config
-def Context.configKey (c : Context) : UInt64 := c.keyedConfig.key
+def Context.configKey (c : Context) : UInt64 := c.keyedConfig.key ||| ((c.canUnfold? matches .standard).toUInt64 <<< 40)
 
 /--
 The `MetaM` monad is a core component of Lean's metaprogramming framework, facilitating the
@@ -1172,7 +1183,10 @@ def elimMVarDeps (xs : Array Expr) (e : Expr) (preserveOrder : Bool := false) :
     { ctx with keyedConfig := { config } }
 
 @[inline] def withCanUnfoldPred (p : Config → ConstantInfo → CoreM Bool) : n α → n α :=
-  mapMetaM <| withReader (fun ctx => { ctx with canUnfold? := p })
+  mapMetaM <| withReader (fun ctx => { ctx with canUnfold? := .customUncached p })
+
+@[inline] def withCanUnfoldAtMatcherPred (p : Config → ConstantInfo → CoreM Bool) : n α → n α :=
+  mapMetaM <| withReader (fun ctx => { ctx with canUnfold? := .atMatcher p })
 
 @[inline] def withIncSynthPending : n α → n α :=
   mapMetaM <| withReader (fun ctx => { ctx with synthPendingDepth := ctx.synthPendingDepth + 1 })
@@ -1267,7 +1281,7 @@ def withTrackingZetaDeltaSet (s : FVarIdSet) : n α → n α :=
 @[inline] private def Context.setTransparency (ctx : Context) (transparency : TransparencyMode) : Context :=
   let config := { ctx.config with transparency }
   -- Recall that `transparency` is stored in the first 3 bits (it has 5 values).
-  let key : UInt64 := ((ctx.configKey >>> (3 : UInt64)) <<< 3) ||| transparency.toUInt64
+  let key : UInt64 := ((ctx.keyedConfig.key >>> (3 : UInt64)) <<< 3) ||| transparency.toUInt64
   { ctx with keyedConfig := { config, key } }
 
 @[inline] def withTransparency (mode : TransparencyMode) : n α → n α :=

src/Lean/Meta/ExprDefEq.lean

@@ -2234,7 +2234,7 @@ inductive DefEqCacheKind where
   | permanent -- problem does not have mvars and we are using standard config, we can use one persistent cache.
 
 private def getDefEqCacheKind (t s : Expr) : MetaM DefEqCacheKind := do
-  if t.hasMVar || s.hasMVar || (← read).canUnfold?.isSome then
+  if t.hasMVar || s.hasMVar || (← read).canUnfold? matches .customUncached _ then
     return .transient
   else
     return .permanent

src/Lean/Meta/GetUnfoldableConst.lean

@@ -32,10 +32,10 @@ def canUnfoldDefault (cfg : Config) (info : ConstantInfo) : CoreM Bool := do
 def canUnfold (info : ConstantInfo) : MetaM Bool := do
   let ctx ← read
   let cfg ← getConfig
-  if let some f := ctx.canUnfold? then
-    f cfg info
-  else
-    canUnfoldDefault cfg info
+  match ctx.canUnfold? with
+  | .customUncached f => f cfg info
+  | .standard => canUnfoldDefault cfg info
+  | .atMatcher f => f cfg info
 
 /--
 Look up a constant name, returning the `ConstantInfo`

src/Lean/Meta/Tactic/Cbv/ControlFlow.lean

@@ -294,8 +294,8 @@ public def withCbvOpaqueGuard (x : MetaM α) : MetaM α := do
   withCanUnfoldPred (fun cfg info => do
     if (← isCbvOpaque info.name) then return false
     match prev with
-    | some f => f cfg info
-    | none =>
+    | .customUncached f => f cfg info
+    | .standard =>
       -- Duplicates `canUnfoldDefault` from `Lean.Meta.GetUnfoldableConst` (private).
       match cfg.transparency with
       | .none => return false
@@ -307,6 +307,7 @@ public def withCbvOpaqueGuard (x : MetaM α) : MetaM α := do
         else if status == .instanceReducible && (m == .instances || m == .implicit) then return true
         else if status == .implicitReducible && m == .implicit then return true
         else return false
+    | .atMatcher f => f cfg info
   ) x
 
 builtin_cbv_simproc ↓ simpCbvCond (@cond _ _ _) := simpCond

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

@@ -595,7 +595,7 @@ private def dischargeUsingAssumption? (e : Expr) : SimpM (Option Expr) := do
 partial def dischargeEqnThmHypothesis? (e : Expr) : MetaM (Option Expr) := do
   assert! isEqnThmHypothesis e
   let mvar ← mkFreshExprSyntheticOpaqueMVar e
-  withCanUnfoldPred canUnfoldAtMatcher do
+  withCanUnfoldAtMatcherPred canUnfoldAtMatcher do
     if let .none ← go? mvar.mvarId! then
       instantiateMVars mvar
     else

src/Lean/Meta/WHNF.lean

@@ -526,7 +526,7 @@ private def whnfMatcher (e : Expr) : MetaM Expr := do
      reduced to expose constructors, without bumping the overall transparency level.  -/
   if (← getTransparency) matches .reducible | .instances | .implicit then
     -- Also unfold some default-reducible constants; see `canUnfoldAtMatcher`
-    withCanUnfoldPred canUnfoldAtMatcher do
+    withCanUnfoldAtMatcherPred canUnfoldAtMatcher do
       whnf e
   else
     -- Do NOT use `canUnfoldAtMatcher` here as it does not affect all/default reducibility and inhibits caching (#2564).
@@ -1082,7 +1082,7 @@ def reduceNat? (e : Expr) : MetaM (Option Expr) :=
 @[inline] private def useWHNFCache (e : Expr) : MetaM Bool := do
   -- We cache only closed terms without expr metavars.
   -- Potential refinement: cache if `e` is not stuck at a metavariable
-  if e.hasFVar || e.hasExprMVar || (← read).canUnfold?.isSome then
+  if e.hasFVar || e.hasExprMVar || (← read).canUnfold? matches .customUncached _ then
     return false
   else
     return true