fix: Core.transform API and uses

src/Lean/Compiler/LCNF/ToDecl.lean

@@ -13,8 +13,8 @@ Inline constants tagged with the `[macroInline]` attribute occurring in `e`.
 -/
 def macroInline (e : Expr) : CoreM Expr :=
   Core.transform e fun e => do
-    let .const declName us := e.getAppFn | return .visit e
-    unless hasMacroInlineAttribute (← getEnv) declName do return .visit e
+    let .const declName us := e.getAppFn | return .continue
+    unless hasMacroInlineAttribute (← getEnv) declName do return .continue
     let val ← Core.instantiateValueLevelParams (← getConstInfo declName) us
     return .visit <| val.beta e.getAppArgs
 
@@ -35,11 +35,11 @@ Inline auxiliary `matcher` applications.
 -/
 partial def inlineMatchers (e : Expr) : CoreM Expr :=
   Meta.MetaM.run' <| Meta.transform e fun e => do
-    let .const declName us := e.getAppFn | return .visit e
-    let some info ← Meta.getMatcherInfo? declName | return .visit e
+    let .const declName us := e.getAppFn | return .continue
+    let some info ← Meta.getMatcherInfo? declName | return .continue
     let numArgs := e.getAppNumArgs
     if numArgs > info.arity then
-      return .visit e
+      return .continue
     else if numArgs < info.arity then
       Meta.forallBoundedTelescope (← Meta.inferType e) (info.arity - numArgs) fun xs _ =>
         return .visit (← Meta.mkLambdaFVars xs (mkAppN e xs))
@@ -71,7 +71,7 @@ private def replaceUnsafeRecNames (value : Expr) : CoreM Expr :=
         return .done (.const safeDeclName us)
       else
         return .done e
-    | _ => return .visit e
+    | _ => return .continue
 
 /--
 Convert the given declaration from the Lean environment into `Decl`.

src/Lean/Elab/Inductive.lean

@@ -378,7 +378,7 @@ where
           args := args.set! i param
         return TransformStep.done (mkAppN f args)
       else
-        return TransformStep.visit e
+        return .continue
     transform ctorType (pre := visit)
 
 private def getResultingUniverse : List InductiveType → TermElabM Level

src/Lean/Elab/MutualDef.lean

@@ -118,11 +118,11 @@ See issues #1389 and #875
 -/
 private def cleanupOfNat (type : Expr) : MetaM Expr := do
   Meta.transform type fun e => do
-    if !e.isAppOfArity ``OfNat 2 then return .visit e
+    if !e.isAppOfArity ``OfNat 2 then return .continue
     let arg ← instantiateMVars e.appArg!
-    if !arg.isAppOfArity ``OfNat.ofNat 3 then return .visit e
+    if !arg.isAppOfArity ``OfNat.ofNat 3 then return .continue
     let argArgs := arg.getAppArgs
-    if !argArgs[0]!.isConstOf ``Nat then return .visit e
+    if !argArgs[0]!.isConstOf ``Nat then return .continue
     let eNew := mkApp e.appFn! argArgs[1]!
     return .done eNew
 
@@ -766,13 +766,11 @@ def processDefDeriving (className : Name) (declName : Name) : TermElabM Bool :=
 
 /-- Remove auxiliary match discriminant let-declarations. -/
 def eraseAuxDiscr (e : Expr) : CoreM Expr := do
-  Core.transform e fun e => match e with
-    | Expr.letE n _ v b .. =>
+  Core.transform e fun e => do
+    if let .letE n _ v b .. := e then
       if isAuxDiscrName n then
-        return TransformStep.visit (b.instantiate1 v)
-      else
-        return TransformStep.visit e
-    | e => return TransformStep.visit e
+        return .visit (b.instantiate1 v)
+    return .continue
 
 partial def checkForHiddenUnivLevels (allUserLevelNames : List Name) (preDefs : Array PreDefinition) : TermElabM Unit :=
   unless (← MonadLog.hasErrors) do

src/Lean/Elab/PreDefinition/Main.lean

@@ -73,9 +73,9 @@ private def betaReduceLetRecApps (preDefs : Array PreDefinition) : MetaM (Array
   preDefs.mapM fun preDef => do
     let value ← transform preDef.value fun e => do
       if e.isApp && e.getAppFn.isLambda && e.getAppArgs.all fun arg => arg.getAppFn.isConst && preDefs.any fun preDef => preDef.declName == arg.getAppFn.constName! then
-        return TransformStep.visit e.headBeta
+        return .visit e.headBeta
       else
-        return TransformStep.visit e
+        return .continue
     return { preDef with value }
 
 private def addAsAxioms (preDefs : Array PreDefinition) : TermElabM Unit := do

src/Lean/Elab/PreDefinition/Structural/Preprocess.lean

@@ -26,10 +26,10 @@ private def shouldBetaReduce (e : Expr) (recFnName : Name) : Bool :=
 -/
 def preprocess (e : Expr) (recFnName : Name) : CoreM Expr :=
   Core.transform e
-   fun e => return TransformStep.visit <|
+   fun e =>
      if shouldBetaReduce e recFnName then
-       e.headBeta
+       return .visit e.headBeta
      else
-       e
+       return .continue
 
 end Lean.Elab.Structural

src/Lean/Elab/PreDefinition/WF/Eqns.lean

@@ -94,8 +94,8 @@ where
         let candidate := mkAppN (mkAppN (mkAppN (mkConst declName us) fixedPrefix) args) extraArgs
         trace[Elab.definition.wf] "found nested WF at discr {candidate}"
         if (← withDefault <| isDefEq candidate e) then
-          return TransformStep.visit candidate
-    return TransformStep.visit e
+          return .visit candidate
+    return .continue
 
 /--
   Simplify `match`-expressions when trying to prove equation theorems for a recursive declaration defined using well-founded recursion.

src/Lean/Expr.lean

@@ -1297,7 +1297,7 @@ If `useZeta = true`, then `let`-expressions are visited. That is, it assumes
 that zeta-reduction (aka let-expansion) is going to be used.
 -/
 def isHeadBetaTarget (e : Expr) (useZeta := false) : Bool :=
-  e.getAppFn.isHeadBetaTargetFn useZeta
+  e.isApp && e.getAppFn.isHeadBetaTargetFn useZeta
 
 private def etaExpandedBody : Expr → Nat → Nat → Option Expr
   | app f (bvar j), n+1, i => if j == i then etaExpandedBody f n (i+1) else none

src/Lean/Meta/Coe.lean

@@ -21,19 +21,13 @@ def isCoeDecl (declName : Name) : Bool :=
 /-- Expand coercions occurring in `e` -/
 partial def expandCoe (e : Expr) : MetaM Expr :=
   withReducibleAndInstances do
-    return (← transform e (pre := step))
-where
-  step (e : Expr) : MetaM TransformStep := do
-    let f := e.getAppFn
-    if !f.isConst then
-      return TransformStep.visit e
-    else
-      let declName := f.constName!
-      if isCoeDecl declName then
-        match (← unfoldDefinition? e) with
-        | none   => return TransformStep.visit e
-        | some e => step e.headBeta
-      else
-        return TransformStep.visit e
+    transform e fun e => do
+      let f := e.getAppFn
+      if f.isConst then
+        let declName := f.constName!
+        if isCoeDecl declName then
+          if let some e ← unfoldDefinition? e then
+            return .visit e.headBeta
+      return .continue
 
 end Lean.Meta

src/Lean/Meta/IndPredBelow.lean

@@ -156,7 +156,7 @@ where
       if let some name := e.constName? then
         if let some _ := ctx.typeInfos.findIdx? fun indVal => indVal.name == name then
           modify (· + 1)
-      return TransformStep.visit e
+      return .continue
 
     if cnt > 1 then
       throwError "only arrows are allowed as premises. Multiple recursive occurrences detected:{indentExpr domain}"

src/Lean/Meta/Injective.lean

@@ -28,7 +28,7 @@ def elimOptParam (type : Expr) : CoreM Expr := do
     if e.isAppOfArity  ``optParam 2 then
       return TransformStep.visit (e.getArg! 0)
     else
-      return TransformStep.visit e
+      return .continue
 
 private partial def mkInjectiveTheoremTypeCore? (ctorVal : ConstructorVal) (useEq : Bool) : MetaM (Option Expr) := do
   let us := ctorVal.levelParams.map mkLevelParam

src/Lean/Meta/Match/MatchEqs.lean

@@ -59,7 +59,7 @@ def unfoldNamedPattern (e : Expr) : MetaM Expr := do
     if let some e := isNamedPattern? e then
       if let some eNew ← unfoldDefinition? e then
         return TransformStep.visit eNew
-    return TransformStep.visit e
+    return .continue
   Meta.transform e (pre := visit)
 
 /--
@@ -545,26 +545,23 @@ where
     transform e fun e => do
       if (← isCastEqRec e) then
         return .done (← convertCastEqRec e)
-      else match e.getAppFn with
-        | Expr.fvar fvarId .. =>
-          match (← read).find? fvarId with
-          | some (altNew, numParams, argMask) =>
-            trace[Meta.Match.matchEqs] ">> argMask: {argMask}, e: {e}, {altNew}"
-            let mut newArgs := #[]
-            let argMask := trimFalseTrail argMask
-            unless e.getAppNumArgs ≥ argMask.size do
-              throwError "unexpected occurrence of `match`-expression alternative (aka minor premise) while creating splitter/eliminator theorem for `{matchDeclName}`, minor premise is partially applied{indentExpr e}\npossible solution if you are matching on inductive families: add its indices as additional discriminants"
-            for arg in e.getAppArgs, includeArg in argMask do
-              if includeArg then
-                newArgs := newArgs.push arg
-            let eNew := mkAppN altNew newArgs
-            /- Recall that `numParams` does not include the equalities associated with discriminants of the form `h : discr`. -/
-            let (mvars, _, _) ← forallMetaBoundedTelescope (← inferType eNew) (numParams - newArgs.size) (kind := MetavarKind.syntheticOpaque)
-            modify fun s => s ++ (mvars.map (·.mvarId!))
-            let eNew := mkAppN eNew mvars
-            return TransformStep.done eNew
-          | none => return TransformStep.visit e
-        | _ => return TransformStep.visit e
+      else
+        let Expr.fvar fvarId .. := e.getAppFn | return .continue
+        let some (altNew, numParams, argMask) := (← read).find? fvarId | return .continue
+        trace[Meta.Match.matchEqs] ">> argMask: {argMask}, e: {e}, {altNew}"
+        let mut newArgs := #[]
+        let argMask := trimFalseTrail argMask
+        unless e.getAppNumArgs ≥ argMask.size do
+          throwError "unexpected occurrence of `match`-expression alternative (aka minor premise) while creating splitter/eliminator theorem for `{matchDeclName}`, minor premise is partially applied{indentExpr e}\npossible solution if you are matching on inductive families: add its indices as additional discriminants"
+        for arg in e.getAppArgs, includeArg in argMask do
+          if includeArg then
+            newArgs := newArgs.push arg
+        let eNew := mkAppN altNew newArgs
+        /- Recall that `numParams` does not include the equalities associated with discriminants of the form `h : discr`. -/
+        let (mvars, _, _) ← forallMetaBoundedTelescope (← inferType eNew) (numParams - newArgs.size) (kind := MetavarKind.syntheticOpaque)
+        modify fun s => s ++ (mvars.map (·.mvarId!))
+        let eNew := mkAppN eNew mvars
+        return TransformStep.done eNew
 
   proveSubgoalLoop (mvarId : MVarId) : MetaM Unit := do
     trace[Meta.Match.matchEqs] "proveSubgoalLoop\n{mvarId}"

src/Lean/Meta/Tactic/Delta.lean

@@ -20,8 +20,8 @@ def delta? (e : Expr) (p : Name → Bool := fun _ => true) : CoreM (Option Expr)
 def deltaExpand (e : Expr) (p : Name → Bool) : CoreM Expr :=
   Core.transform e fun e => do
     match (← delta? e p) with
-    | some e' => return TransformStep.visit e'
-    | none    => return TransformStep.visit e
+    | some e' => return .visit e'
+    | none    => return .continue
 
 /--
 Delta expand declarations that satisfy `p` at `mvarId` type.

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

@@ -180,7 +180,7 @@ private partial def dsimp (e : Expr) : M Expr := do
     if let Step.visit r ← rewritePre e (fun _ => pure none) (rflOnly := true) then
       if r.expr != e then
         return .visit r.expr
-    return .visit e
+    return .continue
   let post (e : Expr) : M TransformStep := do
     if let Step.visit r ← rewritePost e (fun _ => pure none) (rflOnly := true) then
       if r.expr != e then

src/Lean/Meta/Tactic/Split.lean

@@ -111,12 +111,12 @@ private partial def generalizeMatchDiscrs (mvarId : MVarId) (matcherDeclName : N
       let rec mkNewTarget (e : Expr) : MetaM Expr := do
         let pre (e : Expr) : MetaM TransformStep := do
           if !e.isAppOf matcherDeclName || e.getAppNumArgs != matcherInfo.arity then
-            return .visit e
-          let some matcherApp ← matchMatcherApp? e | return .visit e
+            return .continue
+          let some matcherApp ← matchMatcherApp? e | return .continue
           for matcherDiscr in matcherApp.discrs, discr in discrs do
             unless matcherDiscr == discr do
               trace[Meta.Tactic.split] "discr mismatch {matcherDiscr} != {discr}"
-              return .visit e
+              return .continue
           let matcherApp := { matcherApp with discrs := discrVars }
           foundRef.set true
           let mut altsNew := #[]

src/Lean/Meta/Transform.lean

@@ -8,17 +8,23 @@ import Lean.Meta.Basic
 namespace Lean
 
 inductive TransformStep where
-  | done  (e : Expr)
-  | visit (e : Expr)
+  | /-- Return expression without visiting any subexpressions. -/
+    done (e : Expr)
+  | /-- Visit expression (which should be different from current expression) instead. -/
+    visit (e : Expr)
+  | /--
+      Continue transformation with the given expression (defaults to current expression).
+      For `pre`, this means visiting the children of the expression.
+      For `post`, this is equivalent to returning `done`. -/
+    continue (e? : Option Expr := none)
 
 namespace Core
 
 /--
-  Tranform the expression `input` using `pre` and `post`.
-  - `pre s` is invoked before visiting the children of subterm 's'. If the result is `TransformStep.visit sNew`, then
-     `sNew` is traversed by transform. If the result is `TransformStep.done sNew`, then `s` is just replaced with `sNew`.
-     In both cases, `sNew` must be definitionally equal to `s`
-  - `post s` is invoked after visiting the children of subterm `s`.
+  Transform the expression `input` using `pre` and `post`.
+  - First `pre` is invoked with the current expression and recursion is continued according to the `TransformStep` result.
+    In all cases, the expression contained in the result, if any, must be definitionally equal to the current expression.
+  - After recursion, if any, `post` is invoked on the resulting expression.
 
   The term `s` in both `pre s` and `post s` may contain loose bound variables. So, this method is not appropriate for
   if one needs to apply operations (e.g., `whnf`, `inferType`) that do not handle loose bound variables.
@@ -28,8 +34,8 @@ namespace Core
 -/
 partial def transform {m} [Monad m] [MonadLiftT CoreM m] [MonadControlT CoreM m]
     (input : Expr)
-    (pre   : Expr → m TransformStep := fun e => return TransformStep.visit e)
-    (post  : Expr → m TransformStep := fun e => return TransformStep.done e)
+    (pre   : Expr → m TransformStep := fun _ => return .continue)
+    (post  : Expr → m TransformStep := fun e => return .done e)
     : m Expr :=
   let _ : STWorld IO.RealWorld m := ⟨⟩
   let _ : MonadLiftT (ST IO.RealWorld) m := { monadLift := fun x => liftM (m := CoreM) (liftM (m := ST IO.RealWorld) x) }
@@ -37,11 +43,15 @@ partial def transform {m} [Monad m] [MonadLiftT CoreM m] [MonadControlT CoreM m]
     checkCache { val := e : ExprStructEq } fun _ => Core.withIncRecDepth do
       let rec visitPost (e : Expr) : MonadCacheT ExprStructEq Expr m Expr := do
         match (← post e) with
-        | TransformStep.done e  => pure e
-        | TransformStep.visit e => visit e
+        | .done e      => pure e
+        | .visit e     => visit e
+        | .continue e? => pure (e?.getD e)
       match (← pre e) with
-      | TransformStep.done e  => pure e
-      | TransformStep.visit e => match e with
+      | .done e  => pure e
+      | .visit e => visitPost (← visit e)
+      | .continue e? =>
+        let e := e?.getD e
+        match e with
         | Expr.forallE _ d b _ => visitPost (e.updateForallE! (← visit d) (← visit b))
         | Expr.lam _ d b _     => visitPost (e.updateLambdaE! (← visit d) (← visit b))
         | Expr.letE _ t v b _  => visitPost (e.updateLet! (← visit t) (← visit v) (← visit b))
@@ -52,7 +62,7 @@ partial def transform {m} [Monad m] [MonadLiftT CoreM m] [MonadControlT CoreM m]
   visit input |>.run
 
 def betaReduce (e : Expr) : CoreM Expr :=
-  transform e (pre := fun e => return TransformStep.visit e.headBeta)
+  transform e (pre := fun e => return if e.isHeadBetaTarget then .visit e.headBeta else .continue)
 
 end Core
 
@@ -63,8 +73,8 @@ namespace Meta
   So, it is safe to use any `MetaM` method at `pre` and `post`. -/
 partial def transform {m} [Monad m] [MonadLiftT MetaM m] [MonadControlT MetaM m] [MonadTrace m] [MonadRef m] [MonadOptions m] [AddMessageContext m]
     (input : Expr)
-    (pre   : Expr → m TransformStep := fun e => return TransformStep.visit e)
-    (post  : Expr → m TransformStep := fun e => return TransformStep.done e)
+    (pre   : Expr → m TransformStep := fun _ => return .continue)
+    (post  : Expr → m TransformStep := fun e => return .done e)
     (usedLetOnly := false)
     : m Expr := do
   let _ : STWorld IO.RealWorld m := ⟨⟩
@@ -73,8 +83,9 @@ partial def transform {m} [Monad m] [MonadLiftT MetaM m] [MonadControlT MetaM m]
     checkCache { val := e : ExprStructEq } fun _ => Meta.withIncRecDepth do
       let rec visitPost (e : Expr) : MonadCacheT ExprStructEq Expr m Expr := do
         match (← post e) with
-        | TransformStep.done e  => pure e
-        | TransformStep.visit e => visit e
+        | .done e      => pure e
+        | .visit e     => visit e
+        | .continue e? => pure (e?.getD e)
       let rec visitLambda (fvars : Array Expr) (e : Expr) : MonadCacheT ExprStructEq Expr m Expr := do
         match e with
         | Expr.lam n d b c =>
@@ -97,8 +108,11 @@ partial def transform {m} [Monad m] [MonadLiftT MetaM m] [MonadControlT MetaM m]
         e.withApp fun f args => do
           visitPost (mkAppN (← visit f) (← args.mapM visit))
       match (← pre e) with
-      | TransformStep.done e  => pure e
-      | TransformStep.visit e => match e with
+      | .done e  => pure e
+      | .visit e => visit e
+      | .continue e? =>
+        let e := e?.getD e
+        match e with
         | Expr.forallE ..    => visitForall #[] e
         | Expr.lam ..        => visitLambda #[] e
         | Expr.letE ..       => visitLet #[] e
@@ -119,7 +133,7 @@ def zetaReduce (e : Expr) : MetaM Expr := do
           return TransformStep.visit value
         else
           return TransformStep.done e
-    | e => return TransformStep.visit e
+    | _ => return .continue
   transform e (pre := pre) (usedLetOnly := true)
 
 /-- Unfold definitions and theorems in `e` that are not in the current environment, but are in `biggerEnv`. -/
@@ -139,7 +153,7 @@ def unfoldDeclsFrom (biggerEnv : Environment) (e : Expr) : CoreM Expr := do
             return TransformStep.done e
         else
           return TransformStep.done e
-      | _ => return TransformStep.visit e
+      | _ => return .continue
     Core.transform e (pre := pre)
 
 def eraseInaccessibleAnnotations (e : Expr) : CoreM Expr :=

tests/lean/conv1.lean

@@ -173,3 +173,10 @@ example (p : Prop) : p := by
 
 example (p : (n : Nat) → Fin n → Prop) (i : Fin 5) (hp : p 5 i) : p 5 j := by
   conv => arg 1
+
+-- repeated `zeta`
+example : let a := 0; let b := a; b = 0 := by
+  intros
+  conv =>
+    zeta
+    trace_state

tests/lean/conv1.lean.expected.out

@@ -82,3 +82,6 @@ conv1.lean:169:10-169:15: error: invalid 'arg' conv tactic, application has only
 conv1.lean:172:10-172:13: error: invalid 'congr' conv tactic, application or implication expected
   p
 conv1.lean:175:10-175:15: error: cannot select argument
+a✝ : Nat := 0
+b✝ : Nat := a✝
+| 0 = 0

tests/lean/structuralEqns.lean

@@ -16,3 +16,10 @@ def foo (xs ys zs : List Nat) : List Nat :=
 
 #eval tst ``foo
 #check foo._unfold
+
+
+def bar (xs ys : List Nat) : List Nat :=
+  match xs ++ [], ys ++ [] with
+  | xs', ys'   => xs' ++ ys'
+
+#print bar  -- should not contain either `let _discr` aux binding

tests/lean/structuralEqns.lean.expected.out

@@ -9,3 +9,7 @@ foo._unfold : ∀ (xs ys zs : List Nat),
         match ys' with
         | [] => [1]
         | x => [2]
+def bar : List Nat → List Nat → List Nat :=
+fun xs ys =>
+  match xs ++ [], ys ++ [] with
+  | xs', ys' => xs' ++ ys'