fix: withTrackingZetaDelta must reset cache (#6381)

This PR fixes a bug in `withTrackingZetaDelta` and
`withTrackingZetaDeltaSet`. The `MetaM` caches need to be reset. See new
test.

src/Lean/Meta/Basic.lean

@@ -1096,39 +1096,42 @@ def elimMVarDeps (xs : Array Expr) (e : Expr) (preserveOrder : Bool := false) :
 @[inline] def withInTypeClassResolution : n α → n α :=
   mapMetaM <| withReader (fun ctx => { ctx with inTypeClassResolution := true })
 
+@[inline] def withFreshCache : n α → n α :=
+  mapMetaM fun x => do
+    let cacheSaved := (← get).cache
+    modify fun s => { s with cache := {} }
+    try
+      x
+    finally
+      modify fun s => { s with cache := cacheSaved }
+
 /--
 Executes `x` tracking zetaDelta reductions `Config.trackZetaDelta := true`
 -/
 @[inline] def withTrackingZetaDelta : n α → n α :=
-  mapMetaM <| withReader (fun ctx => { ctx with trackZetaDelta := true })
-
-def withZetaDeltaSetImp (s : FVarIdSet) (x : MetaM α) : MetaM α := do
-  if s.isEmpty then
-    x
-  else
-    let cacheSaved := (← get).cache
-    modify fun s => { s with cache := {} }
-    try
-      withReader (fun ctx => { ctx with zetaDeltaSet := s }) x
-    finally
-      modify fun s => { s with cache := cacheSaved }
+  mapMetaM fun x =>
+    withFreshCache <| withReader (fun ctx => { ctx with trackZetaDelta := true }) x
 
 /--
 `withZetaDeltaSet s x` executes `x` with `zetaDeltaSet := s`.
 The cache is reset while executing `x` if `s` is not empty.
 -/
 def withZetaDeltaSet (s : FVarIdSet) : n α → n α :=
-  mapMetaM <| withZetaDeltaSetImp s
+  mapMetaM fun x =>
+    if s.isEmpty then
+      x
+    else
+      withFreshCache <| withReader (fun ctx => { ctx with zetaDeltaSet := s }) x
 
 /--
-Similar to `withZetaDeltaSet`, but also enables `withTrackingZetaDelta` if
-`s` is not empty.
+Similar to `withZetaDeltaSet`, but also enables `withTrackingZetaDelta` if `s` is not empty.
 -/
-def withTrackingZetaDeltaSet (s : FVarIdSet) (x : n α) : n α := do
-  if s.isEmpty then
-    x
-  else
-    withZetaDeltaSet s <| withTrackingZetaDelta x
+def withTrackingZetaDeltaSet (s : FVarIdSet) : n α → n α :=
+  mapMetaM fun x =>
+    if s.isEmpty then
+      x
+    else
+      withFreshCache <| withReader (fun ctx => { ctx with zetaDeltaSet := s, trackZetaDelta := true }) x
 
 @[inline] def withoutProofIrrelevance (x : n α) : n α :=
   withConfig (fun cfg => { cfg with proofIrrelevance := false }) x

tests/lean/run/trackZetaDeltaCacheIssue.lean

@@ -0,0 +1,33 @@
+import Lean
+
+open Lean Meta
+
+def g : Nat → String := fun _ => ""
+
+/--
+info: String
+[A]
+-/
+#guard_msgs in
+run_meta do
+  withLetDecl `A (mkSort 1) (← mkArrow (mkConst ``Nat) (mkConst ``String)) fun A => do
+  withLetDecl `g A (mkConst ``g) fun g => do
+    let e := mkApp g (mkNatLit 0)
+    withTrackingZetaDelta do
+      IO.println (← ppExpr (← inferType e))
+      IO.println s!"{← (← getZetaDeltaFVarIds).toList.mapM fun x => ppExpr (mkFVar x)}"
+
+/--
+info: String
+String
+[A]
+-/
+#guard_msgs in
+run_meta do
+  withLetDecl `A (mkSort 1) (← mkArrow (mkConst ``Nat) (mkConst ``String)) fun A => do
+  withLetDecl `g A (mkConst ``g) fun g => do
+    let e := mkApp g (mkNatLit 0)
+    IO.println (← ppExpr (← inferType e))
+    withTrackingZetaDelta do
+      IO.println (← ppExpr (← inferType e))
+      IO.println s!"{← (← getZetaDeltaFVarIds).toList.mapM fun x => ppExpr (mkFVar x)}"