feat: support for tracking which let-decls have been zeta expanded

src/Lean/Meta/Basic.lean

@@ -31,14 +31,14 @@ namespace Lean
 namespace Meta
 
 structure Config :=
-(foApprox           : Bool    := false)
-(ctxApprox          : Bool    := false)
-(quasiPatternApprox : Bool    := false)
+(foApprox           : Bool := false)
+(ctxApprox          : Bool := false)
+(quasiPatternApprox : Bool := false)
 /- When `constApprox` is set to true,
    we solve `?m t =?= c` using
    `?m := fun _ => c`
    when `?m t` is not a higher-order pattern and `c` is not an application as -/
-(constApprox        : Bool    := false)
+(constApprox        : Bool := false)
 /-
   When the following flag is set,
   `isDefEq` throws the exeption `Exeption.isDefEqStuck`
@@ -47,8 +47,12 @@ structure Config :=
   This feature is useful for type class resolution where
   we may want to notify the caller that the TC problem may be solveable
   later after it assigns `?m`. -/
-(isDefEqStuckEx     : Bool    := false)
+(isDefEqStuckEx     : Bool := false)
 (transparency       : TransparencyMode := TransparencyMode.default)
+/- If zetaNonDep == false, then non dependent let-decls are not zeta expanded. -/
+(zetaNonDep         : Bool := true)
+/- When `trackZeta == true`, we store zetaFVarIds all free variables that have been zeta-expanded. -/
+(trackZeta          : Bool := false)
 
 structure ParamInfo :=
 (implicit     : Bool      := false)
@@ -91,9 +95,11 @@ structure PostponedEntry :=
 (rhs       : Level)
 
 structure State :=
-(mctx           : MetavarContext       := {})
-(cache          : Cache                := {})
-(postponed      : PersistentArray PostponedEntry := {})
+(mctx        : MetavarContext := {})
+(cache       : Cache := {})
+(postponed   : PersistentArray PostponedEntry := {})
+/- When `trackZeta == true`, then any let-decl free variable that is zeta expansion performed by `MetaM` is stored in `zetaFVarIds`. -/
+(zetaFVarIds : NameSet := {})
 
 instance State.inhabited : Inhabited State := ⟨{}⟩
 
@@ -165,6 +171,8 @@ def getConfig : m Config := liftMetaM do ctx ← read; pure ctx.config
 def setMCtx (mctx : MetavarContext) : m Unit := liftMetaM $ modify fun s => { s with mctx := mctx }
 @[inline] def modifyMCtx (f : MetavarContext → MetavarContext) : m Unit :=
 liftMetaM $ modify fun s => { s with mctx := f s.mctx }
+def resetZetaFVarIds : m Unit := liftMetaM $ modify fun s => { s with zetaFVarIds := {} }
+def getZetaFVarIds : m NameSet := liftMetaM do s ← get; pure s.zetaFVarIds
 
 def mkWHNFRef : IO (IO.Ref (Expr → MetaM Expr)) :=
 IO.mkRef $ fun _ => throwError "whnf implementation was not set"
@@ -395,6 +403,9 @@ if xs.isEmpty then pure e else liftMkBindingM $ MetavarContext.elimMVarDeps xs e
 @[inline] def withConfig {α} (f : Config → Config) : n α → n α :=
 mapMetaM fun _ => adaptReader (fun (ctx : Context) => { ctx with config := f ctx.config })
 
+@[inline] def withTrackingZeta {α} (x : n α) : n α :=
+withConfig (fun cfg => { cfg with trackZeta := true }) x
+
 @[inline] def withTransparency {α} (mode : TransparencyMode) : n α → n α :=
 mapMetaM fun _ => withConfig (fun config => { config with transparency := mode })
 

src/Lean/Meta/WHNF.lean

@@ -208,9 +208,16 @@ liftM $ getStuckMVarImp? e
 | e@(Expr.letE _ _ _ _ _),  k => k e
 | e@(Expr.fvar fvarId _),   k => do
   decl ← getLocalDecl fvarId;
-  match decl.value? with
-  | none   => pure e
-  | some v => whnfEasyCases v k
+  match decl with
+  | LocalDecl.cdecl _ _ _ _ _        => pure e
+  | LocalDecl.ldecl _ _ _ _ v nonDep => do
+    cfg ← getConfig;
+    if nonDep && !cfg.zetaNonDep then
+      pure e
+    else do
+      when cfg.trackZeta $
+        modify fun s => { s with zetaFVarIds := s.zetaFVarIds.insert fvarId };
+      whnfEasyCases v k
 | e@(Expr.mvar mvarId _),   k => do
   v? ← getExprMVarAssignment? mvarId;
   match v? with

tests/lean/run/meta7.lean

@@ -0,0 +1,36 @@
+import Lean.Meta
+open Lean
+open Lean.Meta
+
+set_option trace.Meta true
+set_option trace.Meta.isDefEq.step false
+set_option trace.Meta.isDefEq.delta false
+set_option trace.Meta.isDefEq.assign false
+
+def print (msg : MessageData) : MetaM Unit :=
+trace! `Meta.debug msg
+
+def check (x : MetaM Bool) : MetaM Unit :=
+unlessM x $ throwError "check failed"
+
+def ex : Nat × Nat :=
+let x  := 10;
+let ty := Nat → Nat;
+let f  : ty := fun x => x;
+let n  := 20;
+let z  := f 10;
+(let y  : { v : Nat // v = n } := ⟨20, rfl⟩; y.1 + n + f x, z + 10)
+
+def tst1 : MetaM Unit := do
+print "----- tst1 -----";
+c ← getConstInfo `ex;
+lambdaTelescope c.value?.get! fun xs body =>
+  withTrackingZeta do
+    Meta.check body;
+    ys ← getZetaFVarIds;
+    let ys := ys.toList.map mkFVar;
+    print ys;
+    check $ pure $ ys.length == 2;
+    pure ()
+
+#eval tst1