feat: smart unfolding support for nested match-expressions

See #445

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

@@ -48,15 +48,15 @@ where
               unless xs.size >= numParams do
                 throwError "unexpected matcher application alternative{indentExpr alt}\nat application{indentExpr e}"
               let altBody ← visit altBody
-              let containsIdRhs := Option.isSome <| altBody.find? fun e => e.isConstOf ``idRhs
-              if !containsIdRhs then
+              let containsSUnfoldMatch := Option.isSome <| altBody.find? fun e => smartUnfoldingMatch? e |>.isSome
+              if !containsSUnfoldMatch then
                 let altBody ← mkLambdaFVars xs[numParams:xs.size] altBody
-                let altBody ← mkIdRhs altBody
+                let altBody ← markSmartUnfoldigMatchAlt altBody
                 mkLambdaFVars xs[0:numParams] altBody
               else
                 mkLambdaFVars xs altBody
             altsNew := altsNew.push altNew
-          return { matcherApp with alts := altsNew }.toExpr
+          return markSmartUnfoldingMatch { matcherApp with alts := altsNew }.toExpr
       | _ => processApp e
     | _ => return e
 

src/Lean/Meta/WHNF.lean

@@ -27,6 +27,20 @@ register_builtin_option smartUnfolding : Bool := {
   descr := "when computing weak head normal form, use auxiliary definition created for functions defined by structural recursion"
 }
 
+/-- Add auxiliary annotation to indicate the `match`-expression `e` must be reduced when performing smart unfolding. -/
+def markSmartUnfoldingMatch (e : Expr) : Expr :=
+  mkAnnotation `sunfoldMatch e
+
+def smartUnfoldingMatch? (e : Expr) : Option Expr :=
+  annotation? `sunfoldMatch e
+
+/-- Add auxiliary annotation to indicate expression `e` (a `match` alternative rhs) was successfully reduced by smart unfolding. -/
+def markSmartUnfoldigMatchAlt (e : Expr) : Expr :=
+  mkAnnotation `sunfoldMatchAlt e
+
+def smartUnfoldingMatchAlt? (e : Expr) : Option Expr :=
+  annotation? `sunfoldMatchAlt e
+
 /- ===========================
    Helper methods
    =========================== -/
@@ -382,25 +396,69 @@ partial def whnfCore (e : Expr) : MetaM Expr :=
       | none => return e
     | _ => unreachable!
 
-mutual
-  /-- Reduce `e` until `idRhs` application is exposed or it gets stuck.
-      This is a helper method for implementing smart unfolding. -/
-  private partial def whnfUntilIdRhs (e : Expr) : MetaM Expr := do
-    let e ← whnfCore e
-    match (← getStuckMVar? e) with
-    | some mvarId =>
+/--
+  Recall that `_sunfold` auxiliary definitions contains the markers: `markSmartUnfoldigMatch` (*) and `markSmartUnfoldigMatchAlt` (**).
+  For example, consider the following definition
+  ```
+  def r (i j : Nat) : Nat :=
+    i +
+      match j with
+      | Nat.zero => 1
+      | Nat.succ j =>
+        i + match j with
+            | Nat.zero => 2
+            | Nat.succ j => r i j
+  ```
+  produces the following `_sunfold` auxiliary definition with the markers
+  ```
+  def r._sunfold (i j : Nat) : Nat :=
+    i +
+      (*) match j with
+      | Nat.zero => (**) 1
+      | Nat.succ j =>
+        i + (*) match j with
+            | Nat.zero => (**) 2
+            | Nat.succ j => (**) r i j
+  ```
+
+  `match` expressions marked with `markSmartUnfoldigMatch` (*) must be reduced, otherwise the resulting term is not definitionally
+   equal to the given expression. The recursion may be interrupted as soon as the annotation `markSmartUnfoldingAlt` (**) is reached.
+
+  For example, the term `r i j.succ.succ` reduces to the definitionally equal term `i + i * r i j`
+-/
+partial def smartUnfoldingReduce? (e : Expr) : MetaM (Option Expr) :=
+  go e |>.run
+where
+  go (e : Expr) : OptionT MetaM Expr := do
+    match e with
+    | Expr.letE n t v b _ => withLetDecl n t (← go v) fun x => do mkLetFVars #[x] (← go b)
+    | Expr.lam .. => lambdaTelescope e fun xs b => do mkLambdaFVars xs (← go b)
+    | Expr.app f a .. => mkApp (← go f) (← go a)
+    | Expr.proj _ _ s _ => e.updateProj! (← go s)
+    | Expr.mdata _ b _  =>
+      if let some m := smartUnfoldingMatch? e then
+        goMatch m
+      else
+        e.updateMData! (← go b)
+    | _ => return e
+
+  goMatch (e : Expr) : OptionT MetaM Expr := do
+    match (← reduceMatcher? e) with
+    | ReduceMatcherResult.reduced e =>
+      if let some alt := smartUnfoldingMatchAlt? e then
+        return alt
+      else
+        go e
+    | ReduceMatcherResult.stuck e' =>
+      let mvarId ← getStuckMVar? e'
       /- Try to "unstuck" by resolving pending TC problems -/
       if (← Meta.synthPending mvarId) then
-        whnfUntilIdRhs e
+        goMatch e
       else
-        return e -- failed because metavariable is blocking reduction
-    | _ =>
-      if isIdRhsApp e then
-        return e -- done
-      else
-        match (← unfoldDefinition? e) with
-        | some e => whnfUntilIdRhs e
-        | none   => pure e -- failed because of symbolic argument
+        failure
+    | _ => failure
+
+mutual
 
   /--
     Auxiliary method for unfolding a class projection when transparency is set to `TransparencyMode.instances`.
@@ -440,12 +498,8 @@ mutual
           if smartUnfolding.get (← getOptions) then
             match (← getConstNoEx? (mkSmartUnfoldingNameFor fInfo.name)) with
             | some fAuxInfo@(ConstantInfo.defnInfo _) =>
-              deltaBetaDefinition fAuxInfo fLvls e.getAppRevArgs (fun _ => pure none) fun e₁ => do
-                let e₂ ← whnfUntilIdRhs e₁
-                if isIdRhsApp e₂ then
-                  return some (extractIdRhs e₂)
-                else
-                  return none
+              deltaBetaDefinition fAuxInfo fLvls e.getAppRevArgs (fun _ => pure none) fun e₁ =>
+                smartUnfoldingReduce? e₁
             | _ =>
               if (← getMatcherInfo? fInfo.name).isSome then
                 -- Recall that `whnfCore` tries to reduce "matcher" applications.