perf: use WHNF cache in whnfMatcher also at sub-default transparency (#14323)

This PR addresses a performance regression from #13895 where
`Lake.Config.ConfigDecl`, and more concretely the realization of
`ConfigType.eq_1`, becomes significantly slower. The underlying issue
uncovered by #13895 is that `whnfMatcher` uses a custom "can unfold?"
predicate, which indirectly disables the WHNF cache because functions
cannot be used as cache keys, when at reducible, instances or implicit
transparency.

This PR splits the "can unfold?" predicate in `Context` into two parts.
`Context` still contains an optional custom predicate that, when used,
disables the cache. But `whnfMatcher` does not use it anymore. `Config`,
the cacheable structure, gets a new Boolean flag indicating whether we
are at a matcher. Depending on this flag, one of two hard-coded
predicates is used.

Author: datokrat

src/Lean/Meta/Basic.lean

@@ -80,6 +80,16 @@ def ProjReductionKind.toUInt64 : ProjReductionKind → UInt64
   | .yesWithDelta => 2
   | .yesWithDeltaI => 3
 
+structure CanUnfoldPredicateConfig where
+  toBool : Bool
+deriving Inhabited, Repr
+
+@[inline, expose, match_pattern]
+def CanUnfoldPredicateConfig.default : CanUnfoldPredicateConfig := ⟨false⟩
+
+@[inline, expose, match_pattern]
+def CanUnfoldPredicateConfig.atMatcher : CanUnfoldPredicateConfig := ⟨true⟩
+
 /--
 Configuration flags for the `MetaM` monad.
 Many of them are used to control the `isDefEq` function that checks whether two terms are definitionally equal or not.
@@ -194,6 +204,11 @@ structure Config where
   When `zeta := true`, then `zetaHave := false` disables zeta reduction of `have` expressions.
   -/
   zetaHave : Bool := true
+  /--
+  A predicate to control whether a constant can be unfolded or not at `whnf`.
+  Ignored if `Context.customCanUnfoldPredicate?` is set.
+  -/
+  canUnfoldPredicateConfig : CanUnfoldPredicateConfig := .default
   deriving Inhabited, Repr
 
 /-- Convert `isDefEq` and `WHNF` relevant parts into a key for caching results -/
@@ -216,7 +231,8 @@ private def Config.toKey (c : Config) : UInt64 :=
   (c.etaStruct.toUInt64 <<< 17) |||
   (c.proj.toUInt64 <<< 19) |||
   (c.zetaHave.toUInt64 <<< 21) |||
-  (c.zetaUnused.toUInt64 <<< 22)
+  (c.zetaUnused.toUInt64 <<< 22) |||
+  (c.canUnfoldPredicateConfig.toBool.toUInt64 <<< 23)
 
 /-- Configuration with key produced by `Config.toKey`. -/
 structure ConfigWithKey where
@@ -230,6 +246,16 @@ instance : Inhabited ConfigWithKey where  -- #9463
 def Config.toConfigWithKey (c : Config) : ConfigWithKey :=
   { config := c }
 
+def ConfigWithKey.withCanUnfoldAtMatcherPred (c : ConfigWithKey) : ConfigWithKey :=
+  { config := { c.config with canUnfoldPredicateConfig := .atMatcher },
+    key :=
+      have : CanUnfoldPredicateConfig.atMatcher.toBool = true := rfl
+      c.key ||| ((1 : UInt64) <<< 23) }
+
+def ConfigWithKey.setTransparency (transparency : TransparencyMode) (c : ConfigWithKey) : ConfigWithKey :=
+  { config := { c.config with transparency }
+    key := ((c.key >>> (3 : UInt64)) <<< 3) ||| transparency.toUInt64 }
+
 /--
 Function parameter information cache.
 -/
@@ -494,25 +520,29 @@ structure Context where
   /-- Not `none` when inside of an `isDefEq` test. See `PostponedEntry`. -/
   defEqCtx?         : Option DefEqContext  := none
   /--
-    Track the number of nested `synthPending` invocations. Nested invocations can happen
-    when the type class resolution invokes `synthPending`.
+  Track the number of nested `synthPending` invocations. Nested invocations can happen
+  when the type class resolution invokes `synthPending`.
 
-    Remark: `synthPending` fails if `synthPendingDepth > maxSynthPendingDepth`.
+  Remark: `synthPending` fails if `synthPendingDepth > maxSynthPendingDepth`.
   -/
   synthPendingDepth : Nat                  := 0
   /--
-    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
+  A predicate to control whether a constant can be unfolded or not at `whnf`.
+  If set, overrides `Config.canUnfoldPredicateConfig`.
+  Note that we do not cache results at `whnf` when `canUnfold?` is not `none`.
+  This field lives outside `Config` because it is not cacheable and its type does not have a `Repr`
+  instance.
+  -/
+  customCanUnfoldPredicate? : Option (Config → ConstantInfo → CoreM Bool) := none
   /--
   When `Config.univApprox := true`, this flag is set to `true` when there is no
   progress processing universe constraints.
   -/
   univApprox        : Bool := false
   /--
   `inTypeClassResolution := true` when `isDefEq` is invoked at `tryResolve` in the type class
-   resolution module. We don't use `isDefEqProjDelta` when performing TC resolution due to performance issues.
-   This is not a great solution, but a proper solution would require a more sophisticated caching mechanism.
+  resolution module. We don't use `isDefEqProjDelta` when performing TC resolution due to performance issues.
+  This is not a great solution, but a proper solution would require a more sophisticated caching mechanism.
   -/
   inTypeClassResolution : Bool := false
   /--
@@ -1172,7 +1202,12 @@ 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 customCanUnfoldPredicate? := some p })
+
+@[inline] def withCanUnfoldAtMatcherPred : n α → n α :=
+  mapMetaM <| withReader (fun ctx => { ctx with
+    customCanUnfoldPredicate? := none,
+    keyedConfig := ctx.keyedConfig.withCanUnfoldAtMatcherPred })
 
 @[inline] def withIncSynthPending : n α → n α :=
   mapMetaM <| withReader (fun ctx => { ctx with synthPendingDepth := ctx.synthPendingDepth + 1 })
@@ -1264,11 +1299,9 @@ def withTrackingZetaDeltaSet (s : FVarIdSet) : n α → n α :=
 @[inline] def withoutProofIrrelevance (x : n α) : n α :=
   withConfig (fun cfg => { cfg with proofIrrelevance := false }) x
 
-@[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
-  { ctx with keyedConfig := { config, key } }
+@[inline]
+private def Context.setTransparency (ctx : Context) (transparency : TransparencyMode) : Context :=
+  { ctx with keyedConfig := ctx.keyedConfig.setTransparency transparency }
 
 @[inline] def withTransparency (mode : TransparencyMode) : n α → n α :=
   -- We avoid `withConfig` for performance reasons.

src/Lean/Meta/ExprDefEq.lean

@@ -2236,7 +2236,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).customCanUnfoldPredicate?.isSome then
     return .transient
   else
     return .permanent

src/Lean/Meta/GetUnfoldableConst.lean

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 module
 prelude
 public import Lean.Meta.Basic
+public import Lean.Meta.Match.MatchPatternAttr
 public section
 namespace Lean.Meta
 
@@ -29,13 +30,42 @@ def canUnfoldDefault (cfg : Config) (info : ConstantInfo) : CoreM Bool := do
     else
       return false
 
+/--
+Alternative can-unfold predicate used inside `whnfMatcher`.
+See module comment above `unfoldNestedDIte` in `Lean.Meta.WHNF`.
+-/
+def canUnfoldAtMatcher (cfg : Config) (info : ConstantInfo) : CoreM Bool := do
+  if (← canUnfoldDefault cfg info) then
+    return true
+  /- Beyond what the normal transparency allows, we additionally unfold
+     certain definitions to expose constructors in match discriminants. -/
+  if hasMatchPatternAttribute (← getEnv) info.name then
+    return true
+  return info.name == ``OfNat.ofNat -- needed to reduce numeric literals in match discriminants
+   || info.name == ``NatCast.natCast -- needed for `↑m` in match discriminants (pervasive in Int proofs)
+   || info.name == ``Zero.zero || info.name == ``One.one -- needed for `0`/`1` in match discriminants
+   || info.name == ``decEq
+   || info.name == ``Nat.decEq
+   || info.name == ``Char.ofNat   || info.name == ``Char.ofNatAux
+   || info.name == ``String.decEq || info.name == ``List.hasDecEq
+   || info.name == ``Fin.ofNat -- needed for Fin literal reduction in match discriminants
+   || info.name == ``UInt8.ofNat  || info.name == ``UInt8.decEq
+   || info.name == ``UInt16.ofNat || info.name == ``UInt16.decEq
+   || info.name == ``UInt32.ofNat || info.name == ``UInt32.decEq
+   || info.name == ``UInt64.ofNat || info.name == ``UInt64.decEq
+   /- `Fin.ofNat` reduces to `⟨a % n, _⟩`, so we also need to unfold `%` (i.e., `HMod.hMod`
+      and `Mod.mod`) to expose the `Fin.mk` constructor in match discriminants. -/
+   || info.name == ``HMod.hMod || info.name == ``Mod.mod
+
 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.customCanUnfoldPredicate? with
+  | some f => f cfg info
+  | none =>
+    match ctx.config.canUnfoldPredicateConfig with
+    | .default => canUnfoldDefault cfg info
+    | .atMatcher => canUnfoldAtMatcher cfg info
 
 /--
 Look up a constant name, returning the `ConstantInfo`

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

@@ -290,23 +290,16 @@ This prevents kernel-level reduction (used by `reduceRecMatcher?` and `reducePro
 from bypassing the `@[cbv_opaque]` attribute.
 -/
 public def withCbvOpaqueGuard (x : MetaM α) : MetaM α := do
-  let prev := (← readThe Meta.Context).canUnfold?
+  let prevCustomCanUnfoldPredicate? := (← readThe Meta.Context).customCanUnfoldPredicate?
+  let prevCanUnfoldPredicateConfig := (← readThe Meta.Context).config.canUnfoldPredicateConfig
   withCanUnfoldPred (fun cfg info => do
     if (← isCbvOpaque info.name) then return false
-    match prev with
-    | some f => f cfg info
-    | none =>
-      -- Duplicates `canUnfoldDefault` from `Lean.Meta.GetUnfoldableConst` (private).
-      match cfg.transparency with
-      | .none => return false
-      | .all  => return true
-      | .default => return !(← isIrreducible info.name)
-      | m =>
-        let status ← getReducibilityStatus info.name
-        if status == .reducible then return true
-        else if status == .instanceReducible && (m == .instances || m == .implicit) then return true
-        else if status == .implicitReducible && m == .implicit then return true
-        else return false
+    match prevCustomCanUnfoldPredicate? with
+    | .some f => f cfg info
+    | .none =>
+      match prevCanUnfoldPredicateConfig with
+      | .default => canUnfoldDefault cfg info
+      | .atMatcher => canUnfoldAtMatcher 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 do
     if let .none ← go? mvar.mvarId! then
       instantiateMVars mvar
     else

src/Lean/Meta/WHNF.lean

@@ -491,33 +491,6 @@ private def unfoldNestedDIte (e : Expr) : CoreM Expr := do
   else
     return e
 
-/--
-Auxiliary predicate for `whnfMatcher`.
-See comment above.
--/
-def canUnfoldAtMatcher (cfg : Config) (info : ConstantInfo) : CoreM Bool := do
-  if (← canUnfoldDefault cfg info) then
-    return true
-  /- Beyond what the normal transparency allows, we additionally unfold
-     certain definitions to expose constructors in match discriminants. -/
-  if hasMatchPatternAttribute (← getEnv) info.name then
-    return true
-  return info.name == ``OfNat.ofNat -- needed to reduce numeric literals in match discriminants
-   || info.name == ``NatCast.natCast -- needed for `↑m` in match discriminants (pervasive in Int proofs)
-   || info.name == ``Zero.zero || info.name == ``One.one -- needed for `0`/`1` in match discriminants
-   || info.name == ``decEq
-   || info.name == ``Nat.decEq
-   || info.name == ``Char.ofNat   || info.name == ``Char.ofNatAux
-   || info.name == ``String.decEq || info.name == ``List.hasDecEq
-   || info.name == ``Fin.ofNat -- needed for Fin literal reduction in match discriminants
-   || info.name == ``UInt8.ofNat  || info.name == ``UInt8.decEq
-   || info.name == ``UInt16.ofNat || info.name == ``UInt16.decEq
-   || info.name == ``UInt32.ofNat || info.name == ``UInt32.decEq
-   || info.name == ``UInt64.ofNat || info.name == ``UInt64.decEq
-   /- `Fin.ofNat` reduces to `⟨a % n, _⟩`, so we also need to unfold `%` (i.e., `HMod.hMod`
-      and `Mod.mod`) to expose the `Fin.mk` constructor in match discriminants. -/
-   || info.name == ``HMod.hMod || info.name == ``Mod.mod
-
 private def whnfMatcher (e : Expr) : MetaM Expr := do
   /- When reducing `match` expressions at `.reducible`, `.instances` or `.implicit` transparency,
      we use a custom `canUnfoldAtMatcher` predicate that additionally allows unfolding
@@ -526,7 +499,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 do
       whnf e
   else
     -- Do NOT use `canUnfoldAtMatcher` here as it does not affect all/default reducibility and inhibits caching (#2564).
@@ -1082,7 +1055,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).customCanUnfoldPredicate?.isSome then
     return false
   else
     return true

tests/elab/keyedConfig.lean

@@ -0,0 +1,70 @@
+module
+
+import all Lean.Meta.Basic
+
+/-!
+This file proves that optimized bit-fiddling operations on WHNF config keys are correct.
+-/
+
+open Lean Meta
+
+/-!
+# Preparations
+-/
+
+/--
+A term shifted left by `k` (with `3 ≤ k < 64`) is unchanged by `>>> 3 <<< 3`
+(its bits all sit at position ≥ 3, so clearing the low 3 bits is a no-op).
+-/
+theorem shiftLeft_shiftRight_shiftLeft_of_three_le (x k : UInt64) (hk3 : 3 ≤ k.toNat) (hk64 : k.toNat < 64) :
+    (x <<< k) >>> (3:UInt64) <<< 3 = x <<< k := by
+  apply UInt64.toBitVec_inj.1
+  rw [BitVec.eq_of_getLsbD_eq_iff]
+  intro i hi
+  have hkn : (k.toBitVec % 64).toNat = k.toNat := by
+    rw [BitVec.toNat_umod]; exact Nat.mod_eq_of_lt hk64
+  simp only [UInt64.toBitVec_shiftLeft, UInt64.toBitVec_shiftRight,
+             BitVec.shiftLeft_eq', BitVec.ushiftRight_eq',
+             BitVec.getLsbD_shiftLeft, BitVec.getLsbD_ushiftRight,
+             show (UInt64.toBitVec 3 % 64).toNat = 3 from by decide, hkn]
+  rw [show 3 + (i - 3) - k.toNat = i - k.toNat from by omega]
+  rcases Nat.lt_or_ge i 3 with h|h <;> rcases Nat.lt_or_ge i k.toNat with h'|h' <;>
+    simp_all <;> omega
+
+/-!
+# Proofs
+-/
+
+example (c : ConfigWithKey) (h : c.config.toKey = c.key) :
+    let c' := c.withCanUnfoldAtMatcherPred
+    c'.config.toKey = c'.key := by
+  obtain ⟨cfg, k⟩ := c
+  simp only [ConfigWithKey.withCanUnfoldAtMatcherPred] at h ⊢
+  subst h
+  simp only [Config.toKey, CanUnfoldPredicateConfig.atMatcher]
+  cases cfg.canUnfoldPredicateConfig.toBool <;>
+    simp [Bool.toUInt64, UInt64.or_assoc, UInt64.or_self, UInt64.or_zero]
+
+example (c : ConfigWithKey) (t : TransparencyMode) (h : c.config.toKey = c.key) :
+    let c' := c.setTransparency t
+    c'.config.toKey = c'.key := by
+  obtain ⟨cfg, k⟩ := c
+  simp only [ConfigWithKey.setTransparency] at h ⊢
+  subst h
+  simp only [Config.toKey, UInt64.shiftRight_or, UInt64.shiftLeft_or,
+             show cfg.transparency.toUInt64 >>> (3:UInt64) = 0 from by cases cfg.transparency <;> rfl,
+             UInt64.zero_or]
+  repeat rw [shiftLeft_shiftRight_shiftLeft_of_three_le _ _ (by decide) (by decide)]
+  -- Without generalizing these away, the proof is too slow.
+  -- We could use `bv_decide` at the cost of heavier imports.
+  generalize (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a,
+    (_ : UInt64) <<< _ = a, (_ : UInt64) <<< _ = a
+  simp only [UInt64.or_assoc, UInt64.or_comm]