test: add new tests

src/Init/SimpLemmas.lean

@@ -72,6 +72,21 @@ theorem let_body_congr {α : Sort u} {β : α → Sort v} {b b' : (a : α) →
     (a : α) (h : ∀ x, b x = b' x) : (let x := a; b x) = (let x := a; b' x) :=
   (funext h : b = b') ▸ rfl
 
+theorem letFun_unused {α : Sort u} {β : Sort v} (a : α) {b b' : β} (h : b = b') : @letFun α (fun _ => β) a (fun _ => b) = b' :=
+  h
+
+theorem letFun_congr {α : Sort u} {β : Sort v}  {a a' : α} {f f' : α → β} (h₁ : a = a') (h₂ : ∀ x, f x = f' x)
+    : @letFun α (fun _ => β) a f = @letFun α (fun _ => β) a' f' := by
+  rw [h₁, funext h₂]
+
+theorem letFun_body_congr {α : Sort u} {β : Sort v}  (a : α) {f f' : α → β} (h : ∀ x, f x = f' x)
+    : @letFun α (fun _ => β) a f = @letFun α (fun _ => β) a f' := by
+  rw [funext h]
+
+theorem letFun_val_congr {α : Sort u} {β : Sort v} {a a' : α} {f : α → β} (h : a = a')
+    : @letFun α (fun _ => β) a f = @letFun α (fun _ => β) a' f := by
+  rw [h]
+
 @[congr]
 theorem ite_congr {x y u v : α} {s : Decidable b} [Decidable c]
     (h₁ : b = c) (h₂ : c → x = u) (h₃ : ¬ c → y = v) : ite b x y = ite c u v := by

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

@@ -571,10 +571,127 @@ def congr (e : Expr) : SimpM Result := do
   else
     congrDefault e
 
+/--
+Returns `true` if `e` is of the form `@letFun _ (fun _ => β) _ _`
+
+`β` must not contain loose bound variables. Recall that `simp` does not have support for `let_fun`s
+where resulting type depends on the `let`-value. Example:
+```
+let_fun n := 10;
+BitVec.zero n
+```
+-/
+def isNonDepLetFun (e : Expr) : Bool :=
+  let_expr letFun _ beta _ body := e | false
+  beta.isLambda && !beta.bindingBody!.hasLooseBVars && body.isLambda
+
+/--
+Auxiliary structure used to represent the return value of `simpNonDepLetFun.go`.
+-/
+private structure SimpLetFunResult where
+  /--
+  The simplified expression. Note that is may contain loose bound variables.
+  `simpNonDepLetFun.go` attempts to minimize the quadratic overhead imposed
+  by the locally nameless discipline in a sequence of `let_fun` declarations.
+  -/
+  expr     : Expr
+  /--
+  The proof that the simplified expression is equal to the input one.
+  It may containt loose bound variables. See `expr` field.
+  -/
+  proof    : Expr
+  /--
+  `modified := false` iff `expr` is structurally equal to the input expression.
+  -/
+  modified : Bool
+
+/--
+Simplifies a sequence of `let_fun` declarations.
+It attempts to minimize the quadratic overhead imposed by
+the locally nameless discipline.
+-/
+partial def simpNonDepLetFun (e : Expr) : SimpM Result := do
+  let rec go (xs : Array Expr) (e : Expr) : SimpM SimpLetFunResult := do
+    /-
+    Helper function applied when `e` is not a `let_fun` or
+    is a non supported `let_fun` (e.g., the resulting type depends on the value).
+    -/
+    let stop : SimpM SimpLetFunResult := do
+      let e := e.instantiateRev xs
+      let r ← simp e
+      return { expr := r.expr.abstract xs, proof := (← r.getProof).abstract xs, modified :=  r.expr != e }
+    let_expr f@letFun alpha betaFun val body := e | stop
+    let us := f.constLevels!
+    let [_, v] := us | stop
+    /-
+    Recall that `let_fun x : α := val; e[x]` is encoded at
+    ```
+    @letFun α (fun x : α => β[x]) val (fun x : α => e[x])
+    ```
+    `betaFun` is `(fun x : α => β[x])`. If `β[x]` does not have loose bound variables then the resulting type
+    does not depend on the value since it does not depend on `x`.
+
+    We also check whether `alpha` does not depend on the previous `let_fun`s in the sequence.
+    This check is just to make the code simpler. It is not common to have a type depending on the value of a previous `let_fun`.
+    -/
+    if alpha.hasLooseBVars || !betaFun.isLambda || !body.isLambda || betaFun.bindingBody!.hasLooseBVars then
+      stop
+    else if !body.bindingBody!.hasLooseBVar 0 then
+      /-
+      Redundant `let_fun`. The simplifier will remove it.
+      Remark: the `hasLooseBVar` check here may introduce a quadratic overhead in the worst case.
+      If observe that in practice, we may use a separate step for removing unused variables.
+
+      Remark: note that we do **not** simplify the value in this case.
+      -/
+      let x := mkConst `__no_used_dummy__ -- dummy value
+      let { expr, proof, .. } ← go (xs.push x) body.bindingBody!
+      let proof := mkApp6 (mkConst ``letFun_unused us) alpha betaFun.bindingBody! val body.bindingBody! expr proof
+      return { expr, proof, modified := true }
+    else
+      let beta    := betaFun.bindingBody!
+      let valInst := val.instantiateRev xs
+      let valResult ← simp valInst
+      withLocalDecl body.bindingName! body.bindingInfo! alpha fun x => do
+        let valIsNew := valResult.expr != valInst
+        let { expr, proof, modified := bodyIsNew } ← go (xs.push x) body.bindingBody!
+        if !valIsNew && !bodyIsNew then
+          /-
+          Value and body were not simplified. We just return `e` and a new refl proof.
+          We must use the low-level `Expr` APIs because `e` may contain loose bound variables.
+          -/
+          let proof := mkApp2 (mkConst ``Eq.refl [v]) beta e
+          return { expr := e, proof, modified := false }
+        else
+          let body' := mkLambda body.bindingName! body.bindingInfo! alpha expr
+          let val'  := valResult.expr.abstract xs
+          let e'    := mkApp4 f alpha betaFun val' body'
+          if valIsNew && bodyIsNew then
+            -- Value and body were simplified
+            let valProof := (← valResult.getProof).abstract xs
+            let proof := mkApp8 (mkConst ``letFun_congr us) alpha beta val val' body body' valProof (mkLambda body.bindingName! body.bindingInfo! alpha proof)
+            return { expr := e', proof, modified := true }
+          else if valIsNew then
+            -- Only the value was simplified.
+            let valProof := (← valResult.getProof).abstract xs
+            let proof := mkApp6 (mkConst ``letFun_val_congr us) alpha beta val val' body valProof
+            return { expr := e', proof, modified := true }
+          else
+            -- Only the body was simplified.
+            let proof := mkApp6 (mkConst ``letFun_body_congr us) alpha beta val body body' (mkLambda body.bindingName! body.bindingInfo! alpha proof)
+            return { expr := e', proof, modified := true }
+  let { expr, proof, modified } ← go #[] e
+  if !modified then
+    return { expr := e }
+  else
+    return { expr, proof? := proof }
+
 def simpApp (e : Expr) : SimpM Result := do
   if isOfNatNatLit e || isOfScientificLit e || isCharLit e then
     -- Recall that we fold "orphan" kernel Nat literals `n` into `OfNat.ofNat n`
     return { expr := e }
+  else if isNonDepLetFun e then
+    simpNonDepLetFun e
   else
     congr e
 

tests/lean/run/simpLetFunIssue.lean

@@ -0,0 +1,105 @@
+example : (λ x => x)
+           =
+           (λ x : Nat =>
+             let_fun foo := λ y => id (id y)
+             foo (0 + x)) := by
+  simp -zeta only [id]
+  guard_target =ₛ
+           (λ x => x)
+           =
+           (λ x : Nat =>
+             let_fun foo := λ y => y
+             foo (0 + x))
+  simp
+
+example (a : Nat) (h : a = b) : (let_fun x := 1*a; 0 + x) = 0 + b := by
+   simp -zeta only [Nat.zero_add]
+   guard_target =ₛ
+       (let_fun x := 1 * a; x) = b
+   simp -zeta only [Nat.one_mul]
+   guard_target =ₛ
+       (let_fun x := a; x) = b
+   simp [h]
+
+example (a : Nat) (h : a = b) : (let_fun x := 1*a; 0 + x) = 0 + b := by
+   simp -zeta only [Nat.zero_add, Nat.one_mul]
+   guard_target =ₛ
+       (let_fun x := a; x) = b
+   simp [h]
+
+example (a : Nat) (h : a = b) : (let_fun _y := 0; let_fun x := 1*a; 0 + x) = 0 + b := by
+   simp -zeta only [Nat.zero_add, Nat.one_mul]
+   guard_target =ₛ
+       (let_fun x := a; x) = b
+   simp [h]
+
+example (a : Nat) (h : a = b) : (let_fun y := 0; let_fun x := y*0 + 1*a; 0 + x) = 0 + b := by
+   simp -zeta only [Nat.zero_add, Nat.one_mul, Nat.mul_zero]
+   guard_target =ₛ
+       (let_fun x := a; x) = b
+   simp [h]
+
+example (a : Nat) (h : a = b) : (let_fun y := 0; let_fun x := y*0 + 1*a; 0 + x) = 0 + b := by
+   simp -zeta only [Nat.zero_add, Nat.one_mul]
+   guard_target =ₛ
+       (let_fun y := 0; let_fun x := y*0 + a; x) = b
+   fail_if_success simp -zeta only [Nat.zero_add, Nat.one_mul] -- Should not make progress
+   simp -zeta only [Nat.mul_zero]
+   guard_target =ₛ
+       (let_fun x := 0 + a; x) = b
+   simp -zeta only [Nat.zero_add]
+   guard_target =ₛ
+       (let_fun x := a; x) = b
+   simp [h]
+
+def f (n : Nat) (e : Nat) :=
+  match n with
+  | 0 => e
+  | n+1 => let_fun _y := true; let_fun x := 1*e; f n x
+
+example (a b : Nat) (h : a = b) : f 2 (0 + a) = b := by
+   simp -zeta only [f]
+   guard_target =ₛ
+     (let_fun x := 1 * (0 + a); let_fun x := 1 * x; x) = b
+   fail_if_success simp -zeta only [f]
+   simp -zeta only [Nat.one_mul]
+   guard_target =ₛ
+     (let_fun x := 0 + a; let_fun x := x; x) = b
+   simp [h]
+
+example (a b : Nat) (h : a = b) : f 20 (0 + a) = b := by
+   simp -zeta only [f]
+   fail_if_success simp -zeta only [f]
+   simp -zeta only [Nat.one_mul]
+   simp [h]
+
+example (a b : Nat) (h : a = b) : f 50 (0 + a) = b := by
+   simp -zeta only [f]
+   fail_if_success simp -zeta only [f]
+   simp -zeta only [Nat.one_mul]
+   simp [h]
+
+def g (n : Nat) (b : Bool) (e : Nat) :=
+  match n with
+  | 0 => if b then e else 0
+  | n+1 => let_fun b' := !b; let_fun x := 1*e; g n b' x
+
+example (a b : Nat) (h : a = b) : g 2 true (0 + a) = b := by
+  simp -zeta only [g]
+  guard_target =ₛ
+    (let_fun b' := !true; let_fun x := 1 * (0 + a); let_fun b' := !b';
+     let_fun x := 1 * x; if b' = true then x else 0) = b
+  simp -zeta only [Bool.not_true, Nat.one_mul]
+  guard_target =ₛ
+    (let_fun b' := false; let_fun x := 0 + a; let_fun b' := !b';
+     let_fun x := x; if b' = true then x else 0) = b
+  simp [h]
+
+example (a : Nat) : g 33 true (0 + a) = 0 := by
+  simp -zeta only [g]
+  fail_if_success simp -zeta only [g]
+  simp -zeta only [Bool.not_true, Nat.one_mul]
+  fail_if_success simp -zeta only [Bool.not_true, Nat.one_mul]
+  simp -zeta only [Nat.zero_add]
+  fail_if_success simp -zeta only [Nat.zero_add]
+  simp