From c871f66cfaf2ff2dce0aeee96cd8733d381657b4 Mon Sep 17 00:00:00 2001 From: Leonardo de Moura Date: Tue, 6 Jan 2026 11:20:25 -0800 Subject: [PATCH] refactor: `have` telescope support (#11914) This PR factors out the `have`-telescope support used in `simp`, and implements it using the `MonadSimp` interface. The goal is to use this nice infrastructure for both `Meta.simp` and `Sym.simp`. --- src/Lean/Meta.lean | 2 + src/Lean/Meta/HaveTelescope.lean | 411 ++++++++++++++++++++++++++++ src/Lean/Meta/MonadSimp.lean | 25 ++ src/Lean/Meta/Tactic/Simp/Main.lean | 410 ++------------------------- 4 files changed, 459 insertions(+), 389 deletions(-) create mode 100644 src/Lean/Meta/HaveTelescope.lean create mode 100644 src/Lean/Meta/MonadSimp.lean diff --git a/src/Lean/Meta.lean b/src/Lean/Meta.lean index f0c3cb7524..2f491cda05 100644 --- a/src/Lean/Meta.lean +++ b/src/Lean/Meta.lean @@ -59,3 +59,5 @@ public import Lean.Meta.Hint public import Lean.Meta.MethodSpecs public import Lean.Meta.CtorIdxHInj public import Lean.Meta.Sym +public import Lean.Meta.MonadSimp +public import Lean.Meta.HaveTelescope diff --git a/src/Lean/Meta/HaveTelescope.lean b/src/Lean/Meta/HaveTelescope.lean new file mode 100644 index 0000000000..bd9d8b1c2e --- /dev/null +++ b/src/Lean/Meta/HaveTelescope.lean @@ -0,0 +1,411 @@ +/- +Copyright (c) 2025 Lean FRO, LLC. All rights reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Kyle Miller, Leonardo de Moura +-/ +module +prelude +public import Lean.Meta.Basic +public import Lean.Meta.MonadSimp +import Lean.Util.CollectFVars +import Lean.Util.CollectLooseBVars +import Lean.Meta.InferType +import Lean.Meta.AppBuilder +public section +namespace Lean.Meta +/-! +Support for representing `HaveTelescope`s and simplifying them. +We use this module to implement both `Sym.simp` and `Meta.simp` using the `MonadSimp` adapter. +-/ + +/-- +Information about a single `have` in a `have` telescope. +Created by `getHaveTelescopeInfo`. +-/ +structure HaveInfo where + /-- Previous `have`s that the type of this `have` depends on, as indices into `HaveTelescopeInfo.haveInfo`. + Used in `computeFixedUsed` to compute used `have`s. -/ + typeBackDeps : Std.HashSet Nat + /-- Previous `have`s that the value of this `have` depends on, as indices into `HaveTelescopeInfo.haveInfo`. + Used in `computeFixedUsed` to compute used `have`s. -/ + valueBackDeps : Std.HashSet Nat + /-- The local decl for the `have`, so that the local context can be re-entered `have`-by-`have`. + N.B. This stores the fvarid for the `have` as well as cached versions of the value and type + instantiated with the fvars from the telescope. -/ + decl : LocalDecl + /-- The level of the type for this `have`, cached. -/ + level : Level + deriving Inhabited + +/-- +Information about a `have` telescope. +Created by `getHaveTelescopeInfo` and used in `simpHaveTelescope`. + +The data is used to avoid applying `Expr.abstract` more than once on any given subexpression +when constructing terms and proofs during simplification. Abstraction is linear in the size `t` of a term, +and so in a depth-`n` telescope it is `O(nt)`, quadratic in `n`, since `n ≤ t`. +For example, given `have x₁ := v₁; ...; have xₙ := vₙ; b` and `h : b = b'`, we need to construct +```lean +have_body_congr (fun x₁ => ... have_body_congr (fun xₙ => h))) +``` +We apply `Expr.abstract` to `h` *once* and then build the term, rather than +using `mkLambdaFVars #[fvarᵢ] pfᵢ` at each step. + +As an additional optimization, we save the fvar-instantiated terms calculated by `getHaveTelescopeInfo` +so that we don't have to compute them again. This is only saving a constant factor. + +It is also used for computing used `have`s in `computeFixedUsed`. +In `have x₁ := v; have x₂ := x₁; ⋯; have xₙ := xₙ₋₁; b`, if `xₙ` is unused in `b`, then all the +`have`s are unused. Without a global computation of used `have`s, the proof term would be quadratic +in the number of `have`s (with `n` iterations of `simp`). Knowing which `have`s are transitively +unused lets the proof term be linear in size. +-/ +structure HaveTelescopeInfo where + /-- Information about each `have` in the telescope. -/ + haveInfo : Array HaveInfo := #[] + /-- The set of `have`s that the body depends on, as indices into `haveInfo`. + Used in `computeFixedUsed` to compute used `have`s. -/ + bodyDeps : Std.HashSet Nat := {} + /-- The set of `have`s that the type of the body depends on, as indices into `haveInfo`. + This is the set of fixed `have`s. -/ + bodyTypeDeps : Std.HashSet Nat := {} + /-- A cached version of the telescope body, instantiated with fvars from each `HaveInfo.decl`. -/ + body : Expr := Expr.const `_have_telescope_info_dummy_ [] + /-- A cached version of the telescope body's type, instantiated with fvars from each `HaveInfo.decl`. -/ + bodyType : Expr := Expr.const `_have_telescope_info_dummy_ [] + /-- The universe level for the `have` expression, cached. -/ + level : Level := Level.param `_have_telescope_info_dummy_ + deriving Inhabited + +/-- +Efficiently collect dependency information for a `have` telescope. +This is part of a scheme to avoid quadratic overhead from the locally nameless discipline +(see `HaveTelescopeInfo` and `simpHaveTelescope`). + +The expression `e` must not have loose bvars. +-/ +def getHaveTelescopeInfo (e : Expr) : MetaM HaveTelescopeInfo := do + collect e 0 {} (← getLCtx) #[] +where + collect (e : Expr) (numHaves : Nat) (info : HaveTelescopeInfo) (lctx : LocalContext) (fvars : Array Expr) : MetaM HaveTelescopeInfo := do + /- + Gives indices into `fvars` that the uninstantiated expression `a` depends on, from collecting its bvars. + This is more efficient than collecting `fvars` from an instantiated expression, + since the expression likely contains many fvars for the pre-existing local context. + -/ + let getBackDeps (a : Expr) : Std.HashSet Nat := + a.collectLooseBVars.fold (init := {}) fun deps vidx => + -- Since the original expression has no bvars, `vidx < numHaves` must be true. + deps.insert (numHaves - vidx - 1) + match e with + | .letE n t v b true => + let typeBackDeps := getBackDeps t + let valueBackDeps := getBackDeps v + let t := t.instantiateRev fvars + let v := v.instantiateRev fvars + let level ← withLCtx' lctx <| getLevel t + let fvarId ← mkFreshFVarId + let decl := LocalDecl.ldecl 0 fvarId n t v true .default + let info := { info with haveInfo := info.haveInfo.push { typeBackDeps, valueBackDeps, decl, level } } + let lctx := lctx.addDecl decl + let fvars := fvars.push (mkFVar fvarId) + collect b (numHaves + 1) info lctx fvars + | _ => + let bodyDeps := getBackDeps e + withLCtx' lctx do + let body := e.instantiateRev fvars + let bodyType ← inferType body + let level ← getLevel bodyType + -- Collect fvars appearing in the type of `e`. Computing `bodyType` in particular is where `MetaM` is necessary. + let bodyTypeFVarIds := (collectFVars {} bodyType).fvarSet + let bodyTypeDeps : Std.HashSet Nat := Nat.fold fvars.size (init := {}) fun idx _ deps => + if bodyTypeFVarIds.contains fvars[idx].fvarId! then + deps.insert idx + else + deps + return { info with bodyDeps, bodyTypeDeps, body, bodyType, level } + +/-- +Computes which `have`s in the telescope are fixed and which are unused. +The length of the unused array may be less than the number of `have`s: use `unused.getD i true`. +-/ +def HaveTelescopeInfo.computeFixedUsed (info : HaveTelescopeInfo) (keepUnused : Bool) : + MetaM (Array Bool × Array Bool) := do + let fixed ← go info.bodyTypeDeps + if keepUnused then + return (fixed, #[]) + else + let used ← go info.bodyDeps + return (fixed, used) +where + updateArrayFromBackDeps (arr : Array Bool) (s : Std.HashSet Nat) : Array Bool := + s.fold (init := arr) fun arr idx => arr.set! idx true + go init : MetaM (Array Bool) := do + let numHaves := info.haveInfo.size + let mut used : Array Bool := Array.replicate numHaves false + -- Initialize `used` with the body's dependencies. + -- There is no need to consider `info.bodyTypeDeps` in this computation. + used := updateArrayFromBackDeps used init + -- For each used `have`, in reverse order, update `used`. + for i in *...numHaves do + let idx := numHaves - i - 1 + if used[idx]! then + let hinfo := info.haveInfo[idx]! + used := updateArrayFromBackDeps used hinfo.typeBackDeps + used := updateArrayFromBackDeps used hinfo.valueBackDeps + return used + +/-- +Auxiliary structure used to represent the return value of `simpHaveTelescopeAux`. +See `simpHaveTelescope` for an overview and `HaveTelescopeInfo` for an example. +-/ +private structure SimpHaveResult where + /-- + The simplified expression in `(fun x => b) v` form. Note that it may contain loose bound variables. + `simpHaveTelescope` attempts to minimize the quadratic overhead imposed + by the locally nameless discipline in a sequence of `have` expressions. + -/ + expr : Expr + /-- + The type of `expr`. It may contain loose bound variables like in the `expr` field. + Used in proof construction. + -/ + exprType : Expr + /-- + The initial expression in `(fun x => b) v` form. + -/ + exprInit : Expr + /-- + The expression `expr` in `have x := v; b` form. + -/ + exprResult : Expr + /-- + The proof that the simplified expression is equal to the input one. + It may contain loose bound variables like in the `expr` field. + -/ + proof : Expr + /-- + `modified := false` iff `expr` is structurally equal to the input expression. + When false, `proof` is `Eq.refl`. + -/ + modified : Bool + + /- + Re-enters the telescope recorded in `info` using `withExistingLocalDecls` while simplifying according to `fixed`/`used`. + Note that we must use the low-level `Expr` APIs because the expressions may contain loose bound variables. + Note also that we don't enter the body's local context all at once, since we need to be sure that + when we simplify values they have their correct local context. + -/ + deriving Inhabited + +variable {m} [Monad m] [MonadLiftT MetaM m] [MonadControlT MetaM m] [MonadSimp m] + +/- +Re-enters the telescope recorded in `info` using `withExistingLocalDecls` while simplifying according to `fixed`/`used`. +Note that we must use the low-level `Expr` APIs because the expressions may contain loose bound variables. +Note also that we don't enter the body's local context all at once, since we need to be sure that +when we simplify values they have their correct local context. +-/ +private def simpHaveTelescopeAux (info : HaveTelescopeInfo) (fixed : Array Bool) (used : Array Bool) + (e : Expr) (i : Nat) (xs : Array Expr) : m SimpHaveResult := do + if h : i < info.haveInfo.size then + let hinfo := info.haveInfo[i] + -- `x` and `val` are the fvar and value with respect to the local context. + let x := hinfo.decl.toExpr + let val := hinfo.decl.value true + -- Invariant: `v == val.abstract xs`. + let .letE n t v b true := e | unreachable! + -- Universe levels to use for each of the `have_*` theorems. It's the level of `val` and the level of the body. + let us := [hinfo.level, info.level] + if !used.getD i true then + /- + Unused `have`: do not add `x` to the local context then `simp` only the body. + -/ + (do trace[Debug.Meta.Tactic.simp] "have telescope; unused {n} := {val}" : MetaM _) + /- + Complication: Unused `have`s might only be transitively unused. + This means that `b.hasLooseBVar 0` might actually be true. + This matters because `t` and `v` appear in the proof term. + We use a dummy `x` for debugging purposes. (Recall that `Expr.abstract` skips non-fvar/mvars.) + -/ + let x := mkConst `_simp_let_unused_dummy + let rb ← simpHaveTelescopeAux info fixed used b (i + 1) (xs.push x) + -- TODO(kmill): This is a source of quadratic complexity. It might be possible to avoid this, + -- but we will need to carefully re-review the logic (note that `rb.proof` still refers to `x`). + let expr := rb.expr.lowerLooseBVars 1 1 + let exprType := rb.exprType.lowerLooseBVars 1 1 + let exprInit := mkApp (mkLambda n .default t rb.exprInit) v + let exprResult := rb.exprResult.lowerLooseBVars 1 1 + if rb.modified then + let proof := mkApp6 (mkConst ``have_unused_dep' us) t exprType v (mkLambda n .default t rb.exprInit) expr + (mkLambda n .default t rb.proof) + return { expr, exprType, exprInit, exprResult, proof, modified := true } + else + -- If not modified, this must have been a non-transitively unused `have`, so no need for `dep` form. + let proof := mkApp6 (mkConst ``have_unused' us) t exprType v expr expr + (mkApp2 (mkConst ``Eq.refl [info.level]) exprType expr) + return { expr, exprType, exprInit, exprResult, proof, modified := true } + else if fixed.getD i true then + /- + Fixed `have` (like `CongrArgKind.fixed`): dsimp the value and simp the body. + The variable appears in the type of the body. + -/ + let val' ← MonadSimp.dsimp val + (do trace[Debug.Meta.Tactic.simp] "have telescope; fixed {n} := {val} => {val'}" : MetaM _) + let vModified := val != val' + let v' := if vModified then val'.abstract xs else v + withExistingLocalDecls [hinfo.decl] <| MonadSimp.withNewLemmas #[x] do + let rb ← simpHaveTelescopeAux info fixed used b (i + 1) (xs.push x) + let expr := mkApp (mkLambda n .default t rb.expr) v' + let exprType := mkApp (mkLambda n .default t rb.exprType) v' + let exprInit := mkApp (mkLambda n .default t rb.exprInit) v + let exprResult := mkHave n t v' rb.exprResult + -- Note: if `vModified`, then the kernel will reduce the telescopes and potentially do an expensive defeq check. + -- If this is a performance issue, we could try using a `letFun` encoding here make use of the `is_def_eq_args` optimization. + -- Namely, to guide the defeqs via the sequence + -- `(fun x => b) v` = `letFun (fun x => b) v` = `letFun (fun x => b) v'` = `(fun x => b) v'` + if rb.modified then + let proof := mkApp6 (mkConst ``have_body_congr_dep' us) t (mkLambda n .default t rb.exprType) v' + (mkLambda n .default t rb.exprInit) (mkLambda n .default t rb.expr) (mkLambda n .default t rb.proof) + return { expr, exprType, exprInit, exprResult, proof, modified := true } + else + let proof := mkApp2 (mkConst ``Eq.refl [info.level]) exprType expr + return { expr, exprType, exprInit, exprResult, proof, modified := vModified } + else + /- + Non-fixed `have` (like `CongrArgKind.eq`): simp both the value and the body. + The variable does not appear in the type of the body. + -/ + let (v', valModified, vproof) ← match (← MonadSimp.simp val) with + | .rfl => pure (v, false, mkApp2 (mkConst `Eq.refl [hinfo.level]) t v) + | .step val' h => + (do trace[Debug.Meta.Tactic.simp] "have telescope; non-fixed {n} := {val} => {val'}" : MetaM _) + pure (val'.abstract xs, true, h.abstract xs) + withExistingLocalDecls [hinfo.decl] <| MonadSimp.withNewLemmas #[x] do + let rb ← simpHaveTelescopeAux info fixed used b (i + 1) (xs.push x) + let expr := mkApp (mkLambda n .default t rb.expr) v' + assert! !rb.exprType.hasLooseBVar 0 + let exprType := rb.exprType.lowerLooseBVars 1 1 + let exprInit := mkApp (mkLambda n .default t rb.exprInit) v + let exprResult := mkHave n t v' rb.exprResult + match valModified, rb.modified with + | false, false => + let proof := mkApp2 (mkConst `Eq.refl [info.level]) exprType expr + return { expr, exprType, exprInit, exprResult, proof, modified := false } + | true, false => + let proof := mkApp6 (mkConst ``have_val_congr' us) t exprType v v' + (mkLambda n .default t rb.exprInit) vproof + return { expr, exprType, exprInit, exprResult, proof, modified := true } + | false, true => + let proof := mkApp6 (mkConst ``have_body_congr' us) t exprType v + (mkLambda n .default t rb.exprInit) (mkLambda n .default t rb.expr) (mkLambda n .default t rb.proof) + return { expr, exprType, exprInit, exprResult, proof, modified := true } + | true, true => + let proof := mkApp8 (mkConst ``have_congr' us) t exprType v v' + (mkLambda n .default t rb.exprInit) (mkLambda n .default t rb.expr) vproof (mkLambda n .default t rb.proof) + return { expr, exprType, exprInit, exprResult, proof, modified := true } + else + -- Base case: simplify the body. + (do trace[Debug.Meta.Tactic.simp] "have telescope; simplifying body {info.body}" : MetaM _) + let exprType := info.bodyType.abstract xs + match (← MonadSimp.simp info.body) with + | .rfl => + let proof := mkApp2 (mkConst `Eq.refl [info.level]) exprType e + return { expr := e, exprType, exprInit := e, exprResult := e, proof, modified := false } + | .step body' h => + let expr := body'.abstract xs + let proof := h.abstract xs + -- Optimization: if the proof is a `simpHaveTelescope` proof, then remove the type hint + -- to avoid the zeta/beta reductions that the kernel would otherwise do. + -- In `SimpHaveResult.toResult` we insert two type hints; the inner one + -- encodes the `exprInit` and `expr`. + let detectSimpHaveLemma (proof : Expr) : Option SimpHaveResult := do + let_expr id eq proof' := proof | failure + guard <| eq.isAppOfArity ``Eq 3 + let_expr id eq' proof'' := proof' | failure + let_expr Eq _ lhs rhs := eq' | failure + let .const n _ := proof''.getAppFn | failure + guard (n matches ``have_unused_dep' | ``have_unused' | ``have_body_congr_dep' | ``have_val_congr' | ``have_body_congr' | ``have_congr') + return { expr := rhs, exprType, exprInit := lhs, exprResult := expr, proof := proof'', modified := true } + if let some res := detectSimpHaveLemma proof then + return res + return { expr, exprType, exprInit := e, exprResult := expr, proof, modified := true } + +private def SimpHaveResult.toResult (u : Level) (source : Expr) : SimpHaveResult → MonadSimp.Result + | { expr, exprType, exprInit, exprResult, proof, modified, .. } => + if modified then + let proof := + -- Add a type hint to convert back into `have` form. + mkExpectedPropHint (expectedProp := mkApp3 (mkConst ``Eq [u]) exprType source exprResult) <| + -- Add in a second type hint, for use in an optimization to avoid zeta/beta reductions in the kernel + -- The base case in `simpHaveTelescopeAux` detects this construction and uses `exprType`/`exprInit` + -- to construct the `SimpHaveResult`. + -- If the kernel were to support `have` forms of the congruence lemmas this would not be necessary. + mkExpectedPropHint (expectedProp := mkApp3 (mkConst ``Eq [u]) exprType exprInit expr) proof + .step exprResult proof + else + .rfl + +/-- +Routine for simplifying `have` telescopes. Used by `simpLet`. +This is optimized to be able to handle deep `have` telescopes while avoiding quadratic complexity +arising from the locally nameless expression representations. + +## Overview + +Consider a `have` telescope: +``` +have x₁ : T₁ := v₁; ...; have xₙ : Tₙ := vₙ; b +``` +We say `xᵢ` is *fixed* if it appears in the type of `b`. +If `xᵢ` is fixed, then it can only be rewritten using definitional equalities. +Otherwise, we can freely rewrite the value using a propositional equality `vᵢ = vᵢ'`. +The body `b` can always be freely rewritten using a propositional equality `b = b'`. +(All the mentioned equalities must be type correct with respect to the obvious local contexts.) + +If `xᵢ` neither appears in `b` nor any `Tⱼ` nor any `vⱼ`, then its `have` is *unused*. +Unused `have`s can be eliminated, which we do if `cfg.zetaUnused` is true. +We clear `have`s from the end, to be able to transitively clear chains of unused `have`s. +This is why we honor `zetaUnused` here even though `reduceStep` is also responsible for it; +`reduceStep` can only eliminate unused `have`s at the start of a telescope. +Eliminating all transitively unused `have`s at once like this also avoids quadratic complexity. + +If `debug.simp.check.have` is enabled then we typecheck the resulting expression and proof. + +## Proof generation + +There are two main complications with generating proofs. +1. We work almost exclusively with expressions with loose bound variables, + to avoid overhead from instantiating and abstracting free variables, + which can lead to complexity quadratic in the telescope length. + (See `HaveTelescopeInfo`.) +2. We want to avoid triggering zeta reductions in the kernel. + +Regarding this second point, the issue is that we cannot organize the proof using congruence theorems +in the obvious way. For example, given +``` +theorem have_congr {α : Sort u} {β : Sort v} {a a' : α} {f f' : α → β} + (h₁ : a = a') (h₂ : ∀ x, f x = f' x) : + (have x := a; f x) = (have x := a'; f' x) +``` +the kernel sees that the type of `have_congr (fun x => b) (fun x => b') h₁ h₂` +is `(have x := a; (fun x => b) x) = (have x := a'; (fun x => b') x)`, +since when instantiating values it does not beta reduce at the same time. +Hence, when `is_def_eq` is applied to the LHS and `have x := a; b` for example, +it will do `whnf_core` to both sides. + +Instead, we use the form `(fun x => b) a = (fun x => b') a'` in the proofs, +which we can reliably match up without triggering beta reductions in the kernel. +The zeta/beta reductions are then limited to the type hint for the entire telescope, +when we convert this back into `have` form. +In the base case, we include an optimization to avoid unnecessary zeta/beta reductions, +by detecting a `simpHaveTelescope` proofs and removing the type hint. +-/ +def simpHaveTelescope (e : Expr) (zetaUnused := true) : m MonadSimp.Result := do + let info ← getHaveTelescopeInfo e + assert! !info.haveInfo.isEmpty + let (fixed, used) ← info.computeFixedUsed (keepUnused := !zetaUnused) + let r ← simpHaveTelescopeAux info fixed used e 0 (Array.mkEmpty info.haveInfo.size) + return r.toResult info.level e + +end Lean.Meta diff --git a/src/Lean/Meta/MonadSimp.lean b/src/Lean/Meta/MonadSimp.lean new file mode 100644 index 0000000000..c590519e11 --- /dev/null +++ b/src/Lean/Meta/MonadSimp.lean @@ -0,0 +1,25 @@ +/- +Copyright (c) 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved. +Released under Apache 2.0 license as described in the file LICENSE. +Authors: Leonardo de Moura +-/ +module +prelude +public import Lean.Expr +public section +namespace Lean.Meta +/-! +Abstract simplifier API +-/ + +inductive MonadSimp.Result where + | rfl + | step (e : Expr) (h : Expr) + deriving Inhabited + +class MonadSimp (m : Type → Type) where + withNewLemmas (xs : Array Expr) (k : m α) : m α + dsimp : Expr → m Expr + simp : Expr → m MonadSimp.Result + +end Lean.Meta diff --git a/src/Lean/Meta/Tactic/Simp/Main.lean b/src/Lean/Meta/Tactic/Simp/Main.lean index f27329176d..553536e153 100644 --- a/src/Lean/Meta/Tactic/Simp/Main.lean +++ b/src/Lean/Meta/Tactic/Simp/Main.lean @@ -4,18 +4,17 @@ Released under Apache 2.0 license as described in the file LICENSE. Authors: Leonardo de Moura -/ module - prelude public import Lean.Meta.Tactic.Replace public import Lean.Meta.Tactic.Simp.Rewrite public import Lean.Meta.Tactic.Simp.Diagnostics public import Lean.Meta.Match.Value +public import Lean.Meta.MonadSimp public import Lean.Util.CollectLooseBVars +import Lean.Meta.HaveTelescope import Lean.PrettyPrinter import Lean.ExtraModUses - public section - namespace Lean.Meta namespace Simp @@ -263,6 +262,16 @@ def withNewLemmas {α} (xs : Array Expr) (f : SimpM α) : SimpM α := do else withFreshCache do f +instance : MonadSimp SimpM where + simp e := do + let r ← simp e + if r.expr == e then + return .rfl + else + return .step r.expr (← r.getProof) + dsimp := dsimp + withNewLemmas := withNewLemmas + def simpProj (e : Expr) : SimpM Result := do match (← withSimpMetaConfig <| reduceProj? e) with | some e => return { expr := e } @@ -378,393 +387,16 @@ def simpForall (e : Expr) : SimpM Result := withParent e do else return { expr := (← dsimp e) } -/-- -Information about a single `have` in a `have` telescope. -Created by `getHaveTelescopeInfo`. --/ -structure HaveInfo where - /-- Previous `have`s that the type of this `have` depends on, as indices into `HaveTelescopeInfo.haveInfo`. - Used in `computeFixedUsed` to compute used `have`s. -/ - typeBackDeps : Std.HashSet Nat - /-- Previous `have`s that the value of this `have` depends on, as indices into `HaveTelescopeInfo.haveInfo`. - Used in `computeFixedUsed` to compute used `have`s. -/ - valueBackDeps : Std.HashSet Nat - /-- The local decl for the `have`, so that the local context can be re-entered `have`-by-`have`. - N.B. This stores the fvarid for the `have` as well as cached versions of the value and type - instantiated with the fvars from the telescope. -/ - decl : LocalDecl - /-- The level of the type for this `have`, cached. -/ - level : Level - deriving Inhabited - -/-- -Information about a `have` telescope. -Created by `getHaveTelescopeInfo` and used in `simpHaveTelescope`. - -The data is used to avoid applying `Expr.abstract` more than once on any given subexpression -when constructing terms and proofs during simplification. Abstraction is linear in the size `t` of a term, -and so in a depth-`n` telescope it is `O(nt)`, quadratic in `n`, since `n ≤ t`. -For example, given `have x₁ := v₁; ...; have xₙ := vₙ; b` and `h : b = b'`, we need to construct -```lean -have_body_congr (fun x₁ => ... have_body_congr (fun xₙ => h))) -``` -We apply `Expr.abstract` to `h` *once* and then build the term, rather than -using `mkLambdaFVars #[fvarᵢ] pfᵢ` at each step. - -As an additional optimization, we save the fvar-instantiated terms calculated by `getHaveTelescopeInfo` -so that we don't have to compute them again. This is only saving a constant factor. - -It is also used for computing used `have`s in `computeFixedUsed`. -In `have x₁ := v; have x₂ := x₁; ⋯; have xₙ := xₙ₋₁; b`, if `xₙ` is unused in `b`, then all the -`have`s are unused. Without a global computation of used `have`s, the proof term would be quadratic -in the number of `have`s (with `n` iterations of `simp`). Knowing which `have`s are transitively -unused lets the proof term be linear in size. --/ -structure HaveTelescopeInfo where - /-- Information about each `have` in the telescope. -/ - haveInfo : Array HaveInfo := #[] - /-- The set of `have`s that the body depends on, as indices into `haveInfo`. - Used in `computeFixedUsed` to compute used `have`s. -/ - bodyDeps : Std.HashSet Nat := {} - /-- The set of `have`s that the type of the body depends on, as indices into `haveInfo`. - This is the set of fixed `have`s. -/ - bodyTypeDeps : Std.HashSet Nat := {} - /-- A cached version of the telescope body, instantiated with fvars from each `HaveInfo.decl`. -/ - body : Expr := Expr.const `_have_telescope_info_dummy_ [] - /-- A cached version of the telescope body's type, instantiated with fvars from each `HaveInfo.decl`. -/ - bodyType : Expr := Expr.const `_have_telescope_info_dummy_ [] - /-- The universe level for the `have` expression, cached. -/ - level : Level := Level.param `_have_telescope_info_dummy_ - deriving Inhabited - -/-- -Efficiently collect dependency information for a `have` telescope. -This is part of a scheme to avoid quadratic overhead from the locally nameless discipline -(see `HaveTelescopeInfo` and `simpHaveTelescope`). - -The expression `e` must not have loose bvars. --/ -def getHaveTelescopeInfo (e : Expr) : MetaM HaveTelescopeInfo := do - collect e 0 {} (← getLCtx) #[] -where - collect (e : Expr) (numHaves : Nat) (info : HaveTelescopeInfo) (lctx : LocalContext) (fvars : Array Expr) : MetaM HaveTelescopeInfo := do - /- - Gives indices into `fvars` that the uninstantiated expression `a` depends on, from collecting its bvars. - This is more efficient than collecting `fvars` from an instantiated expression, - since the expression likely contains many fvars for the pre-existing local context. - -/ - let getBackDeps (a : Expr) : Std.HashSet Nat := - a.collectLooseBVars.fold (init := {}) fun deps vidx => - -- Since the original expression has no bvars, `vidx < numHaves` must be true. - deps.insert (numHaves - vidx - 1) - match e with - | .letE n t v b true => - let typeBackDeps := getBackDeps t - let valueBackDeps := getBackDeps v - let t := t.instantiateRev fvars - let v := v.instantiateRev fvars - let level ← withLCtx' lctx <| getLevel t - let fvarId ← mkFreshFVarId - let decl := LocalDecl.ldecl 0 fvarId n t v true .default - let info := { info with haveInfo := info.haveInfo.push { typeBackDeps, valueBackDeps, decl, level } } - let lctx := lctx.addDecl decl - let fvars := fvars.push (mkFVar fvarId) - collect b (numHaves + 1) info lctx fvars - | _ => - let bodyDeps := getBackDeps e - withLCtx' lctx do - let body := e.instantiateRev fvars - let bodyType ← inferType body - let level ← getLevel bodyType - -- Collect fvars appearing in the type of `e`. Computing `bodyType` in particular is where `MetaM` is necessary. - let bodyTypeFVarIds := (collectFVars {} bodyType).fvarSet - let bodyTypeDeps : Std.HashSet Nat := Nat.fold fvars.size (init := {}) fun idx _ deps => - if bodyTypeFVarIds.contains fvars[idx].fvarId! then - deps.insert idx - else - deps - return { info with bodyDeps, bodyTypeDeps, body, bodyType, level } - -/-- -Computes which `have`s in the telescope are fixed and which are unused. -The length of the unused array may be less than the number of `have`s: use `unused.getD i true`. --/ -def HaveTelescopeInfo.computeFixedUsed (info : HaveTelescopeInfo) (keepUnused : Bool) : - MetaM (Array Bool × Array Bool) := do - let fixed ← go info.bodyTypeDeps - if keepUnused then - return (fixed, #[]) - else - let used ← go info.bodyDeps - return (fixed, used) -where - updateArrayFromBackDeps (arr : Array Bool) (s : Std.HashSet Nat) : Array Bool := - s.fold (init := arr) fun arr idx => arr.set! idx true - go init : MetaM (Array Bool) := do - let numHaves := info.haveInfo.size - let mut used : Array Bool := Array.replicate numHaves false - -- Initialize `used` with the body's dependencies. - -- There is no need to consider `info.bodyTypeDeps` in this computation. - used := updateArrayFromBackDeps used init - -- For each used `have`, in reverse order, update `used`. - for i in *...numHaves do - let idx := numHaves - i - 1 - if used[idx]! then - let hinfo := info.haveInfo[idx]! - used := updateArrayFromBackDeps used hinfo.typeBackDeps - used := updateArrayFromBackDeps used hinfo.valueBackDeps - return used - -/-- -Auxiliary structure used to represent the return value of `simpHaveTelescopeAux`. -See `simpHaveTelescope` for an overview and `HaveTelescopeInfo` for an example. --/ -private structure SimpHaveResult where - /-- - The simplified expression in `(fun x => b) v` form. Note that it may contain loose bound variables. - `simpHaveTelescope` attempts to minimize the quadratic overhead imposed - by the locally nameless discipline in a sequence of `have` expressions. - -/ - expr : Expr - /-- - The type of `expr`. It may contain loose bound variables like in the `expr` field. - Used in proof construction. - -/ - exprType : Expr - /-- - The initial expression in `(fun x => b) v` form. - -/ - exprInit : Expr - /-- - The expression `expr` in `have x := v; b` form. - -/ - exprResult : Expr - /-- - The proof that the simplified expression is equal to the input one. - It may contain loose bound variables like in the `expr` field. - -/ - proof : Expr - /-- - `modified := false` iff `expr` is structurally equal to the input expression. - When false, `proof` is `Eq.refl`. - -/ - modified : Bool - -private def SimpHaveResult.toResult (u : Level) (source : Expr) : SimpHaveResult → Result - | { expr, exprType, exprInit, exprResult, proof, modified, .. } => - { expr := exprResult - proof? := - if modified then - -- Add a type hint to convert back into `have` form. - some <| mkExpectedPropHint (expectedProp := mkApp3 (mkConst ``Eq [u]) exprType source exprResult) <| - -- Add in a second type hint, for use in an optimization to avoid zeta/beta reductions in the kernel - -- The base case in `simpHaveTelescopeAux` detects this construction and uses `exprType`/`exprInit` - -- to construct the `SimpHaveResult`. - -- If the kernel were to support `have` forms of the congruence lemmas this would not be necessary. - mkExpectedPropHint (expectedProp := mkApp3 (mkConst ``Eq [u]) exprType exprInit expr) proof - else - none } - -/-- -Routine for simplifying `have` telescopes. Used by `simpLet`. -This is optimized to be able to handle deep `have` telescopes while avoiding quadratic complexity -arising from the locally nameless expression representations. - -## Overview - -Consider a `have` telescope: -``` -have x₁ : T₁ := v₁; ...; have xₙ : Tₙ := vₙ; b -``` -We say `xᵢ` is *fixed* if it appears in the type of `b`. -If `xᵢ` is fixed, then it can only be rewritten using definitional equalities. -Otherwise, we can freely rewrite the value using a propositional equality `vᵢ = vᵢ'`. -The body `b` can always be freely rewritten using a propositional equality `b = b'`. -(All the mentioned equalities must be type correct with respect to the obvious local contexts.) - -If `xᵢ` neither appears in `b` nor any `Tⱼ` nor any `vⱼ`, then its `have` is *unused*. -Unused `have`s can be eliminated, which we do if `cfg.zetaUnused` is true. -We clear `have`s from the end, to be able to transitively clear chains of unused `have`s. -This is why we honor `zetaUnused` here even though `reduceStep` is also responsible for it; -`reduceStep` can only eliminate unused `have`s at the start of a telescope. -Eliminating all transitively unused `have`s at once like this also avoids quadratic complexity. - -If `debug.simp.check.have` is enabled then we typecheck the resulting expression and proof. - -## Proof generation - -There are two main complications with generating proofs. -1. We work almost exclusively with expressions with loose bound variables, - to avoid overhead from instantiating and abstracting free variables, - which can lead to complexity quadratic in the telescope length. - (See `HaveTelescopeInfo`.) -2. We want to avoid triggering zeta reductions in the kernel. - -Regarding this second point, the issue is that we cannot organize the proof using congruence theorems -in the obvious way. For example, given -``` -theorem have_congr {α : Sort u} {β : Sort v} {a a' : α} {f f' : α → β} - (h₁ : a = a') (h₂ : ∀ x, f x = f' x) : - (have x := a; f x) = (have x := a'; f' x) -``` -the kernel sees that the type of `have_congr (fun x => b) (fun x => b') h₁ h₂` -is `(have x := a; (fun x => b) x) = (have x := a'; (fun x => b') x)`, -since when instantiating values it does not beta reduce at the same time. -Hence, when `is_def_eq` is applied to the LHS and `have x := a; b` for example, -it will do `whnf_core` to both sides. - -Instead, we use the form `(fun x => b) a = (fun x => b') a'` in the proofs, -which we can reliably match up without triggering beta reductions in the kernel. -The zeta/beta reductions are then limited to the type hint for the entire telescope, -when we convert this back into `have` form. -In the base case, we include an optimization to avoid unnecessary zeta/beta reductions, -by detecting a `simpHaveTelescope` proofs and removing the type hint. --/ +/-- Adapter for `Meta.simpHaveTelescope` -/ def simpHaveTelescope (e : Expr) : SimpM Result := do - Prod.fst <$> withTraceNode `Debug.Meta.Tactic.simp (fun - | .ok (_, used, fixed, modified) => pure m!"{checkEmoji} have telescope; used: {used}; fixed: {fixed}; modified: {modified}" - | .error ex => pure m!"{crossEmoji} {ex.toMessageData}") do - let info ← getHaveTelescopeInfo e - assert! !info.haveInfo.isEmpty - let (fixed, used) ← info.computeFixedUsed (keepUnused := !(← getConfig).zetaUnused) - let r ← simpHaveTelescopeAux info fixed used e 0 (Array.mkEmpty info.haveInfo.size) - if r.modified && debug.simp.check.have.get (← getOptions) then - check r.expr - check r.proof - return (r.toResult info.level e, used, fixed, r.modified) -where - /- - Re-enters the telescope recorded in `info` using `withExistingLocalDecls` while simplifying according to `fixed`/`used`. - Note that we must use the low-level `Expr` APIs because the expressions may contain loose bound variables. - Note also that we don't enter the body's local context all at once, since we need to be sure that - when we simplify values they have their correct local context. - -/ - simpHaveTelescopeAux (info : HaveTelescopeInfo) (fixed : Array Bool) (used : Array Bool) (e : Expr) (i : Nat) (xs : Array Expr) : SimpM SimpHaveResult := do - if h : i < info.haveInfo.size then - let hinfo := info.haveInfo[i] - -- `x` and `val` are the fvar and value with respect to the local context. - let x := hinfo.decl.toExpr - let val := hinfo.decl.value true - -- Invariant: `v == val.abstract xs`. - let .letE n t v b true := e | unreachable! - -- Universe levels to use for each of the `have_*` theorems. It's the level of `val` and the level of the body. - let us := [hinfo.level, info.level] - if !used.getD i true then - /- - Unused `have`: do not add `x` to the local context then `simp` only the body. - -/ - trace[Debug.Meta.Tactic.simp] "have telescope; unused {n} := {val}" - /- - Complication: Unused `have`s might only be transitively unused. - This means that `b.hasLooseBVar 0` might actually be true. - This matters because `t` and `v` appear in the proof term. - We use a dummy `x` for debugging purposes. (Recall that `Expr.abstract` skips non-fvar/mvars.) - -/ - let x := mkConst `_simp_let_unused_dummy - let rb ← simpHaveTelescopeAux info fixed used b (i + 1) (xs.push x) - -- TODO(kmill): This is a source of quadratic complexity. It might be possible to avoid this, - -- but we will need to carefully re-review the logic (note that `rb.proof` still refers to `x`). - let expr := rb.expr.lowerLooseBVars 1 1 - let exprType := rb.exprType.lowerLooseBVars 1 1 - let exprInit := mkApp (mkLambda n .default t rb.exprInit) v - let exprResult := rb.exprResult.lowerLooseBVars 1 1 - if rb.modified then - let proof := mkApp6 (mkConst ``have_unused_dep' us) t exprType v (mkLambda n .default t rb.exprInit) expr - (mkLambda n .default t rb.proof) - return { expr, exprType, exprInit, exprResult, proof, modified := true } - else - -- If not modified, this must have been a non-transitively unused `have`, so no need for `dep` form. - let proof := mkApp6 (mkConst ``have_unused' us) t exprType v expr expr - (mkApp2 (mkConst ``Eq.refl [info.level]) exprType expr) - return { expr, exprType, exprInit, exprResult, proof, modified := true } - else if fixed.getD i true then - /- - Fixed `have` (like `CongrArgKind.fixed`): dsimp the value and simp the body. - The variable appears in the type of the body. - -/ - let val' ← dsimp val - trace[Debug.Meta.Tactic.simp] "have telescope; fixed {n} := {val} => {val'}" - let vModified := val != val' - let v' := if vModified then val'.abstract xs else v - withExistingLocalDecls [hinfo.decl] <| withNewLemmas #[x] do - let rb ← simpHaveTelescopeAux info fixed used b (i + 1) (xs.push x) - let expr := mkApp (mkLambda n .default t rb.expr) v' - let exprType := mkApp (mkLambda n .default t rb.exprType) v' - let exprInit := mkApp (mkLambda n .default t rb.exprInit) v - let exprResult := mkHave n t v' rb.exprResult - -- Note: if `vModified`, then the kernel will reduce the telescopes and potentially do an expensive defeq check. - -- If this is a performance issue, we could try using a `letFun` encoding here make use of the `is_def_eq_args` optimization. - -- Namely, to guide the defeqs via the sequence - -- `(fun x => b) v` = `letFun (fun x => b) v` = `letFun (fun x => b) v'` = `(fun x => b) v'` - if rb.modified then - let proof := mkApp6 (mkConst ``have_body_congr_dep' us) t (mkLambda n .default t rb.exprType) v' - (mkLambda n .default t rb.exprInit) (mkLambda n .default t rb.expr) (mkLambda n .default t rb.proof) - return { expr, exprType, exprInit, exprResult, proof, modified := true } - else - let proof := mkApp2 (mkConst ``Eq.refl [info.level]) exprType expr - return { expr, exprType, exprInit, exprResult, proof, modified := vModified } - else - /- - Non-fixed `have` (like `CongrArgKind.eq`): simp both the value and the body. - The variable does not appear in the type of the body. - -/ - let rval' ← simp val - trace[Debug.Meta.Tactic.simp] "have telescope; non-fixed {n} := {val} => {rval'.expr}" - let valModified := rval'.expr != val - let v' := if valModified then rval'.expr.abstract xs else v - let vproof ← if valModified then - pure <| (← rval'.getProof).abstract xs - else - pure <| mkApp2 (mkConst `Eq.refl [hinfo.level]) t v - withExistingLocalDecls [hinfo.decl] <| withNewLemmas #[x] do - let rb ← simpHaveTelescopeAux info fixed used b (i + 1) (xs.push x) - let expr := mkApp (mkLambda n .default t rb.expr) v' - assert! !rb.exprType.hasLooseBVar 0 - let exprType := rb.exprType.lowerLooseBVars 1 1 - let exprInit := mkApp (mkLambda n .default t rb.exprInit) v - let exprResult := mkHave n t v' rb.exprResult - match valModified, rb.modified with - | false, false => - let proof := mkApp2 (mkConst `Eq.refl [info.level]) exprType expr - return { expr, exprType, exprInit, exprResult, proof, modified := false } - | true, false => - let proof := mkApp6 (mkConst ``have_val_congr' us) t exprType v v' - (mkLambda n .default t rb.exprInit) vproof - return { expr, exprType, exprInit, exprResult, proof, modified := true } - | false, true => - let proof := mkApp6 (mkConst ``have_body_congr' us) t exprType v - (mkLambda n .default t rb.exprInit) (mkLambda n .default t rb.expr) (mkLambda n .default t rb.proof) - return { expr, exprType, exprInit, exprResult, proof, modified := true } - | true, true => - let proof := mkApp8 (mkConst ``have_congr' us) t exprType v v' - (mkLambda n .default t rb.exprInit) (mkLambda n .default t rb.expr) vproof (mkLambda n .default t rb.proof) - return { expr, exprType, exprInit, exprResult, proof, modified := true } - else - -- Base case: simplify the body. - trace[Debug.Meta.Tactic.simp] "have telescope; simplifying body {info.body}" - let r ← simp info.body - let exprType := info.bodyType.abstract xs - if r.expr == info.body then - let proof := mkApp2 (mkConst `Eq.refl [info.level]) exprType e - return { expr := e, exprType, exprInit := e, exprResult := e, proof, modified := false } - else - let expr := r.expr.abstract xs - let proof := (← r.getProof).abstract xs - -- Optimization: if the proof is a `simpHaveTelescope` proof, then remove the type hint - -- to avoid the zeta/beta reductions that the kernel would otherwise do. - -- In `SimpHaveResult.toResult` we insert two type hints; the inner one - -- encodes the `exprInit` and `expr`. - let detectSimpHaveLemma (proof : Expr) : Option SimpHaveResult := do - let_expr id eq proof' := proof | failure - guard <| eq.isAppOfArity ``Eq 3 - let_expr id eq' proof'' := proof' | failure - let_expr Eq _ lhs rhs := eq' | failure - let .const n _ := proof''.getAppFn | failure - guard (n matches ``have_unused_dep' | ``have_unused' | ``have_body_congr_dep' | ``have_val_congr' | ``have_body_congr' | ``have_congr') - return { expr := rhs, exprType, exprInit := lhs, exprResult := expr, proof := proof'', modified := true } - if let some res := detectSimpHaveLemma proof then - return res - return { expr, exprType, exprInit := e, exprResult := expr, proof, modified := true } + let zetaUnused := (← getConfig).zetaUnused + match (← Meta.simpHaveTelescope e zetaUnused) with + | .rfl => return { expr := e } + | .step e' h => + if debug.simp.check.have.get (← getOptions) then + check e' + check h + return { expr := e', proof? := h } /-- Routine for simplifying `let` expressions.