From 53acd3e355c0067e0a16cc773850289a798080d1 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Wed, 22 Jun 2022 19:27:04 -0700
Subject: [PATCH] feat: allow `ext` conv tactic to go inside `let`-declarations

---
 src/Lean/Elab/Tactic/Conv/Congr.lean        | 43 ++++++++++++++++++---
 tests/lean/convZetaLetExt.lean              | 24 ++++++++++++
 tests/lean/convZetaLetExt.lean.expected.out | 29 ++++++++++++++
 3 files changed, 90 insertions(+), 6 deletions(-)
 create mode 100644 tests/lean/convZetaLetExt.lean
 create mode 100644 tests/lean/convZetaLetExt.lean.expected.out

diff --git a/src/Lean/Elab/Tactic/Conv/Congr.lean b/src/Lean/Elab/Tactic/Conv/Congr.lean
index e1d6d2611d..c6359585e3 100644
--- a/src/Lean/Elab/Tactic/Conv/Congr.lean
+++ b/src/Lean/Elab/Tactic/Conv/Congr.lean
@@ -106,22 +106,53 @@ private def selectIdx (tacticName : String) (mvarIds : List (Option MVarId)) (i
       selectIdx "arg" mvarIds i
    | _ => throwUnsupportedSyntax
 
+def extLetBodyCongr? (mvarId : MVarId) (lhs rhs : Expr) : MetaM (Option MVarId) := do
+  match lhs with
+  | .letE n t v b _ =>
+    let u₁ ← getLevel t
+    let f := mkLambda n .default t b
+    unless (← isTypeCorrect f) do
+      throwError "failed to abstract let-expression, result is not type correct"
+    let (β, u₂, f') ← withLocalDeclD n t fun a => do
+      let type ← inferType (mkApp f a)
+      let β ← mkLambdaFVars #[a] type
+      let u₂ ← getLevel type
+      let rhsBody ← mkFreshExprMVar type
+      let f' ← mkLambdaFVars #[a] rhsBody
+      let rhs' := mkLet n t v f'.bindingBody!
+      trace[Meta.debug] "rhs: {rhs'}"
+      unless (← isDefEq rhs rhs') do
+        throwError "failed to go inside let-declaration, type error"
+      return (β, u₂, f')
+    let (arg, mvarId') ← withLocalDeclD n t fun x => do
+      let eqLhs := f.beta #[x]
+      let eqRhs := f'.beta #[x]
+      let mvarNew ← mkFreshExprSyntheticOpaqueMVar (← mkEq eqLhs eqRhs)
+      let arg ← mkLambdaFVars #[x] mvarNew
+      return (arg, mvarNew.mvarId!)
+    let val := mkApp6 (mkConst ``let_body_congr [u₁, u₂]) t β f f' v arg
+    assignExprMVar mvarId val
+    return some (← markAsConvGoal mvarId')
+  | _ => return none
+
 private def extCore (mvarId : MVarId) (userName? : Option Name) : MetaM MVarId :=
    withMVarContext mvarId do
      let userNames := if let some userName := userName? then [userName] else []
-     let (lhs, _) ← getLhsRhsCore mvarId
-     let lhs ← instantiateMVars lhs
-     if lhs.isForall then
+     let (lhs, rhs) ← getLhsRhsCore mvarId
+     let lhs := (← instantiateMVars lhs).consumeMDataAndTypeAnnotations
+     let goForall (mvarId : MVarId) := do
        let [mvarId, _] ← apply mvarId (← mkConstWithFreshMVarLevels ``forall_congr) | throwError "'apply forall_congr' unexpected result"
        let (_, mvarId) ← introN mvarId 1 userNames
        markAsConvGoal mvarId
+     if lhs.isForall then
+       goForall mvarId
+     else if let some mvarId ← extLetBodyCongr? mvarId lhs rhs then
+       return mvarId
      else
        let lhsType ← whnfD (← inferType lhs)
        unless lhsType.isForall do
          throwError "invalid 'ext' conv tactic, function or arrow expected{indentD m!"{lhs} : {lhsType}"}"
-       let [mvarId] ← apply mvarId (← mkConstWithFreshMVarLevels ``funext) | throwError "'apply funext' unexpected result"
-       let (_, mvarId) ← introN mvarId 1 userNames
-       markAsConvGoal mvarId
+       goForall mvarId
 
 private def ext (userName? : Option Name) : TacticM Unit := do
   replaceMainGoal [← extCore (← getMainGoal) userName?]
diff --git a/tests/lean/convZetaLetExt.lean b/tests/lean/convZetaLetExt.lean
new file mode 100644
index 0000000000..04fc86ec38
--- /dev/null
+++ b/tests/lean/convZetaLetExt.lean
@@ -0,0 +1,24 @@
+def f (n : Nat) := n + 1
+
+example (k : Nat) (h : let x := 10; f x = k) : 11 = k := by
+  have : f 10 = 11 := rfl
+  conv at h => zeta; rw [this]
+  trace_state
+  exact h
+
+example (k y : Nat) (h : let x := y; f (0+x) = k) : f y = k := by
+  conv at h => ext x; lhs; arg 1; rw [Nat.zero_add]
+  trace_state
+  exact h
+
+example (g : Nat → Nat) (y : Nat) (h : let x := y; g (0+x) = 0+x) : g y = 0 + y := by
+  conv at h => enter [x, 1, 1]; rw [Nat.zero_add]
+  trace_state
+  exact h
+
+example (g : Nat → Nat) (y : Nat) (h : let x := y; let z := y + 1; g (0+x) = 0+z) : g y = y + 1 := by
+  conv at h => enter [x, z, 1, 1]; rw [Nat.zero_add]
+  trace_state
+  conv at h => enter [x, z, 2]; rw [Nat.zero_add]
+  trace_state
+  exact h
diff --git a/tests/lean/convZetaLetExt.lean.expected.out b/tests/lean/convZetaLetExt.lean.expected.out
new file mode 100644
index 0000000000..c88cb4da30
--- /dev/null
+++ b/tests/lean/convZetaLetExt.lean.expected.out
@@ -0,0 +1,29 @@
+k : Nat
+h : 11 = k
+this : f 10 = 11
+⊢ 11 = k
+k y : Nat
+h :
+  let x := y;
+  f x = k
+⊢ f y = k
+g : Nat → Nat
+y : Nat
+h :
+  let x := y;
+  g x = 0 + x
+⊢ g y = 0 + y
+g : Nat → Nat
+y : Nat
+h :
+  let x := y;
+  let z := y + 1;
+  g x = 0 + z
+⊢ g y = y + 1
+g : Nat → Nat
+y : Nat
+h :
+  let x := y;
+  let z := y + 1;
+  g x = z
+⊢ g y = y + 1