fix: preserve expected types of nested proofs

src/Lean/Meta/AbstractNestedProofs.lean

@@ -14,10 +14,9 @@ public section
 
 namespace Lean.Meta
 
-/-- Abstracts the given proof into an auxiliary theorem, suitably pre-processing its type. -/
-def abstractProof [Monad m] [MonadLiftT MetaM m] [MonadEnv m] [MonadOptions m] [MonadFinally m]
-    (proof : Expr) (cache := true) (postprocessType : Expr → m Expr := pure) : m Expr := do
-  let type ← withoutExporting do inferType proof
+private def abstractProofAt [Monad m] [MonadLiftT MetaM m] [MonadEnv m] [MonadOptions m]
+    [MonadFinally m] (proof type : Expr) (cache := true)
+    (postprocessType : Expr → m Expr := pure) : m Expr := do
   let type ← (Core.betaReduce type : MetaM _)
   let type ← zetaReduce type
   let type ← postprocessType type
@@ -30,6 +29,12 @@ def abstractProof [Monad m] [MonadLiftT MetaM m] [MonadEnv m] [MonadOptions m] [
     In a benchmark created by @selsam, The extra `check` step was a bottleneck. -/
   mkAuxTheorem (cache := cache) type proof (zetaDelta := true)
 
+/-- Abstracts the given proof into an auxiliary theorem, suitably pre-processing its type. -/
+def abstractProof [Monad m] [MonadLiftT MetaM m] [MonadEnv m] [MonadOptions m] [MonadFinally m]
+    (proof : Expr) (cache := true) (postprocessType : Expr → m Expr := pure) : m Expr := do
+  let type ← withoutExporting do inferType proof
+  abstractProofAt proof type cache postprocessType
+
 namespace AbstractNestedProofs
 
 def getLambdaBody (e : Expr) : Expr :=
@@ -67,9 +72,9 @@ def isNonTrivialProof (e : Expr) : MetaM Bool := do
 structure Context where
   cache    : Bool
 
-abbrev M := ReaderT Context $ MonadCacheT ExprStructEq Expr MetaM
+abbrev M := ReaderT Context $ MonadCacheT (ExprStructEq × Option ExprStructEq) Expr MetaM
 
-partial def visit (e : Expr) : M Expr := do
+partial def visit (e : Expr) (expectedType? : Option Expr := none) : M Expr := do
   checkSystem "abstract nested proofs"
   if e.isAtomic then
     pure e
@@ -87,22 +92,48 @@ partial def visit (e : Expr) : M Expr := do
            | none       => pure localDecl
         lctx := lctx.modifyLocalDecl xFVarId fun _ => localDecl
       withLCtx lctx localInstances k
-    checkCache { val := e : ExprStructEq } fun _ => do
-      if (← isNonTrivialProof e) && !e.hasSorry then
+    let isNonTrivialProof := (← isNonTrivialProof e) && !e.hasSorry
+    let expectedTypeKey? :=
+      if isNonTrivialProof then expectedType?.map fun type => { val := type : ExprStructEq }
+      else none
+    checkCache ({ val := e : ExprStructEq }, expectedTypeKey?) fun _ => do
+      if isNonTrivialProof then
         /- Ensure proofs nested in type are also abstracted.
            We skip abstraction for proofs containing `sorry` to avoid generating extra
            "declaration uses sorry" warnings for auxiliary theorems: one per abstracted proof
            instead of a single warning for the main declaration. Additionally, the `zetaDelta`
            expansion in `mkAuxTheorem` can inline let-bound sorry values, causing warnings
            even for proofs that only transitively reference sorry-containing definitions. -/
-        abstractProof e (← read).cache visit
-      else match e with
+        let type ← withoutExporting do
+          match expectedType? with
+          | some type => pure type
+          | none => inferType e
+        abstractProofAt e type (← read).cache (fun type => visit type)
+      else
+        match e with
         | .lam ..
         | .letE ..     => lambdaLetTelescope e fun xs b => visitBinders xs do mkLambdaFVars xs (← visit b) (usedLetOnly := false) (generalizeNondepLet := false)
         | .forallE ..  => forallTelescope e fun xs b => visitBinders xs do mkForallFVars xs (← visit b)
         | .mdata _ b   => return e.updateMData! (← visit b)
         | .proj _ _ b  => return e.updateProj! (← visit b)
-        | .app ..      => e.withApp fun f args => return mkAppN (← visit f) (← args.mapM visit)
+        | .app ..      => e.withApp fun f args => do
+          let mut result ← visit f
+          let mut resultType ← inferType result
+          let mut useExpectedType := true
+          for arg in args do
+            let mut argExpectedType? := none
+            let mut body? := none
+            if useExpectedType then
+              match ← whnf resultType with
+              | .forallE _ type body _ =>
+                argExpectedType? := some type
+                body? := some body
+              | _ => useExpectedType := false
+            let arg ← visit arg argExpectedType?
+            if let some body := body? then
+              resultType := body.instantiate1 arg
+            result := .app result arg
+          return result
         | _            => pure e
 
 end AbstractNestedProofs

tests/lean/simpHoistedProof.lean

@@ -0,0 +1,18 @@
+opaque Bytes : Type
+opaque Bytes.length : Bytes → Nat
+
+axiom foo {α : Type} (f : α → Bytes) (len : Nat)
+    (h : ∀ x, (f x).length = len) : α
+
+def f {len : Nat} (b : { b : Bytes // b.length = len }) : Bytes :=
+  b.val
+
+noncomputable def bar (len : Nat) : { b : Bytes // b.length = len } :=
+  foo f len (by simp [f])
+
+set_option linter.tacticCheckInstances true
+
+#guard_msgs in
+example (len : Nat) : bar len = bar len := by
+  unfold bar
+  rfl