fix handling of lets in WF after discussion with Joachim

Author: kmill

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

@@ -85,9 +85,6 @@ where
       withLocalDecl n c (← loop F d) fun x => do
         mkForallFVars #[x] (← loop F (b.instantiate1 x))
     | Expr.letE n type val body nondep =>
-      -- Convert `have`s to `let`s if they aren't propositions so that the values
-      -- of local declarations can be used within decreasing-by proofs.
-      let nondep ← pure nondep <&&> Meta.isProp type
       mapLetDecl n (← loop F type) (← loop F val) (nondep := nondep) (usedLetOnly := false) fun x => do
         loop F (body.instantiate1 x)
     | Expr.mdata d b =>

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

@@ -244,8 +244,6 @@ where
     | Expr.letE n type val body nondep =>
       loop param type
       loop param val
-      -- Convert `have`s to `let`s if they aren't propositions so that the values of local declarations can be used.
-      let nondep ← pure nondep <&&> Meta.isProp type
       withLetDecl n type val (nondep := nondep) fun x => do
         loop param (body.instantiate1 x)
     | Expr.mdata _d b =>

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

@@ -110,16 +110,68 @@ builtin_dsimproc paramMatcher (_) := fun e => do
   let matcherApp' := { matcherApp with discrs := discrs', alts := alts' }
   return .continue <| matcherApp'.toExpr
 
-/-- `let x := (wfParam e); body[x] ==> let x := e; body[wfParam y]`
-Also accepts non-Prop `have`s as well. -/
+private def anyLetValueIsWfParam (e : Expr) : Bool :=
+  match e with
+  | .letE _ _ v b _ => (isWfParam? v).isSome || anyLetValueIsWfParam b
+  | _               => false
+
+private def numLetsWithValueNotIsWfParam (e : Expr) (acc := 0) : Nat :=
+  match e with
+  | .letE _ _ v b _ => if (isWfParam? v).isSome then acc else numLetsWithValueNotIsWfParam b (acc + 1)
+  | _               => acc
+
+private partial def processParamLet (e : Expr) : MetaM Expr := do
+  if let .letE _ t v b _ := e then
+    if let some v := isWfParam? v then
+      if ← Meta.isProp t then
+        processParamLet <| e.updateLetE! t v b
+      else
+        let u ← getLevel t
+        let b' := b.instantiate1 <| mkApp2 (.const ``wfParam [u]) t (.bvar 0)
+        processParamLet <| e.updateLetE! t v b'
+    else
+      let num := numLetsWithValueNotIsWfParam e
+      assert! num > 0
+      letBoundedTelescope e num fun xs b => do
+        let b' ← processParamLet b
+        mkLetFVars (usedLetOnly := false) (generalizeNondepLet := false) xs b'
+  else
+    return e
+
+/--
+`let x : T := (wfParam e); body[x] ==> let x : T := e; body[wfParam y]` if `T` is not a proposition,
+otherwise `... ==> let x : T := e; body[x]`. (Applies to `have`s too.)
+
+Note: simprocs are provided the head of a let telescope, but not intermediate lets.
+-/
 builtin_dsimproc paramLet (_) := fun e => do
-  let .letE _ t v b _ := e | return .continue
-  let some v := isWfParam? v | return .continue
-  if ← Meta.isProp t then return .continue <| e.updateLetE! t v b
-  let u ← getLevel t
-  let body' := b.instantiate1 <|
-    mkApp2 (.const ``wfParam [u]) t (.bvar 0)
-  return .continue <| e.updateLetE! t v body'
+  unless e.isLet || anyLetValueIsWfParam e do return .continue
+  return .continue (← processParamLet e)
+
+/--
+Transforms non-Prop `have`s to `let`s, so that the values can be used in GuessLex and decreasing-by proofs.
+These `have`s may have been introdued by `simp`, which converts `let`s to `have`s.
+-/
+private def nonPropHaveToLet (e : Expr) : MetaM Expr := do
+  Meta.transform e (pre := fun e => do
+    if (← Meta.isProof e) then
+      return .done e
+    else if e.isLet then
+      -- Recall that `Meta.transform` processes entire let telescopes,
+      -- so we need to handle the telescope all at once.
+      let lctx ← getLCtx
+      let e' ← letTelescope e fun xs b => do
+        let lctx' ← xs.foldlM (init := lctx) fun lctx' x => do
+          let decl ← x.fvarId!.getDecl
+          -- Clear the flag if it's not a prop.
+          let decl' := decl.setNondep <| ← pure decl.isNondep <&&> Meta.isProp decl.type
+          pure <| lctx'.addDecl decl'
+        withLCtx' lctx' do
+          mkLetFVars (usedLetOnly := false) (generalizeNondepLet := false) xs b
+      return .continue e'
+    else
+      return .continue
+  )
 
 def preprocess (e : Expr) : MetaM Simp.Result := do
   unless wf.preprocess.get (← getOptions) do
@@ -143,9 +195,13 @@ def preprocess (e : Expr) : MetaM Simp.Result := do
           if h : as.size ≥ 2 then
             return .continue (mkAppN as[1] as[2:])
         return .continue
+
+    -- Transform `have`s to `let`s for non-propositions.
+    let e'' ← nonPropHaveToLet e''
+
     let result := { result with expr := e'' }
 
-    trace[Elab.definition.wf] "Attach-introduction:{indentExpr e'}\nto{indentExpr result.expr}\ncleaned up as{indentExpr e''}"
+    trace[Elab.definition.wf] "Attach-introduction:{indentExpr e'}\nto{indentExpr result.expr}"
     result.addLambdas xs
 
 end Lean.Elab.WF

src/Lean/LocalContext.lean

@@ -176,6 +176,11 @@ def setNondep : LocalDecl → Bool → LocalDecl
   | ldecl idx id n t v _ k, nd => ldecl idx id n t v nd k
   | d, _                       => d
 
+/-- Returns `true` if this is an `ldecl` with `nondep := true`. -/
+def isNondep : LocalDecl → Bool
+  | ldecl (nondep := nondep) .. => nondep
+  | _                           => false
+
 def setUserName : LocalDecl → Name → LocalDecl
   | cdecl index id _ type bi k,     userName => cdecl index id userName type bi k
   | ldecl index id _ type val nd k, userName => ldecl index id userName type val nd k

tests/lean/run/wf_preprocess.lean

@@ -19,6 +19,27 @@ def Tree.map (f : α → β) (t : Tree α) : Tree β :=
 termination_by t
 decreasing_by trace_state; cases t; decreasing_tactic
 
+/-!
+Checking that the attaches make their way through `let`s.
+-/
+/--
+trace: α : Type u_1
+t : Tree α
+cs : List (Tree α) := t.cs
+t' : Tree α
+h✝ : t' ∈ cs
+⊢ sizeOf t' < sizeOf t
+-/
+#guard_msgs(trace) in
+def Tree.map' (f : α → β) (t : Tree α) : Tree β :=
+  have n := 22
+  let v := t.val
+  let cs := t.cs
+  have : n = n := rfl
+  ⟨f v, cs.map (fun t' => t'.map' f)⟩
+termination_by t
+decreasing_by trace_state; cases t; decreasing_tactic
+
 /--
 info: equations:
 theorem Tree.map.eq_1.{u_1, u_2} : ∀ {α : Type u_1} {β : Type u_2} (f : α → β) (t : Tree α),