lean4-htt/src/Lean/Meta/FunInfo.lean
Kim Morrison 21e8a99eff
feat: refactor of find functions on List/Array/Vector (#6833)
This PR makes the signatures of `find` functions across
`List`/`Array`/`Vector` consistent. Verification lemmas will follow in
subsequent PRs.

We were previously quite inconsistent about the signature of
`indexOf`/`findIdx` functions across `List` and `Array`. Moreover, there
are still quite large gaps in the verification lemma coverage for these
even at the `List` level.

My intention is to make the signatures consistent by providing:
`findIdx` / `findIdx?` / `findFinIdx?` (these all take a predicate, and
return respectively a `Nat`, `Option Nat`, `Option (Fin l.length)`) and
similarly `idxOf` / `idxOf?` / `finIdxOf?` (which look for an element)
for each of List/Array/Vector. I've seen enough examples by now where
each variant is genuinely the most convenient at the call-site, so I'm
going to accept the cost of having many closely related functions.
*Hopefully* for the verification lemmas we can simp all of these into
"projections" of the `Option (Fin l.length)` versions, and then only
have to specify that.

However, I will not plan on immediately either filling in the missing
verification lemmas (or even deciding what the simp normal forms
relating these operations are), and just reach parity amongst
List/Array/Vector for what is already there.
2025-01-30 01:14:21 +00:00

102 lines
4.2 KiB
Text
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

/-
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
prelude
import Lean.Meta.Basic
import Lean.Meta.InferType
namespace Lean.Meta
@[inline] private def checkFunInfoCache (fn : Expr) (maxArgs? : Option Nat) (k : MetaM FunInfo) : MetaM FunInfo := do
let key ← mkInfoCacheKey fn maxArgs?
match (← get).cache.funInfo.find? key with
| some finfo => return finfo
| none => do
let finfo ← k
modify fun s => { s with cache := { s.cache with funInfo := s.cache.funInfo.insert key finfo } }
return finfo
@[inline] private def whenHasVar {α} (e : Expr) (deps : α) (k : αα) : α :=
if e.hasFVar then k deps else deps
private def collectDeps (fvars : Array Expr) (e : Expr) : Array Nat :=
let rec visit (e : Expr) (deps : Array Nat) : Array Nat :=
match e with
| .app f a => whenHasVar e deps (visit a ∘ visit f)
| .forallE _ d b _ => whenHasVar e deps (visit b ∘ visit d)
| .lam _ d b _ => whenHasVar e deps (visit b ∘ visit d)
| .letE _ t v b _ => whenHasVar e deps (visit b ∘ visit v ∘ visit t)
| .proj _ _ e => visit e deps
| .mdata _ e => visit e deps
| .fvar .. =>
match fvars.idxOf? e with
| none => deps
| some i => if deps.contains i then deps else deps.push i
| _ => deps
let deps := visit e #[]
deps.qsort (fun i j => i < j)
/-- Update `hasFwdDeps` fields using new `backDeps` -/
private def updateHasFwdDeps (pinfo : Array ParamInfo) (backDeps : Array Nat) : Array ParamInfo :=
if backDeps.size == 0 then
pinfo
else
-- update hasFwdDeps fields
pinfo.mapIdx fun i info =>
if info.hasFwdDeps then
info
else if backDeps.contains i then
{ info with hasFwdDeps := true }
else
info
private def getFunInfoAux (fn : Expr) (maxArgs? : Option Nat) : MetaM FunInfo :=
checkFunInfoCache fn maxArgs? do
let fnType ← inferType fn
withAtLeastTransparency TransparencyMode.default do
forallBoundedTelescope fnType maxArgs? fun fvars type => do
let mut paramInfo := #[]
let mut higherOrderOutParams : FVarIdSet := {}
for h : i in [:fvars.size] do
let fvar := fvars[i]
let decl ← getFVarLocalDecl fvar
let backDeps := collectDeps fvars decl.type
let dependsOnHigherOrderOutParam :=
!higherOrderOutParams.isEmpty
&& Option.isSome (decl.type.find? fun e => e.isFVar && higherOrderOutParams.contains e.fvarId!)
paramInfo := updateHasFwdDeps paramInfo backDeps
paramInfo := paramInfo.push {
backDeps, dependsOnHigherOrderOutParam
binderInfo := decl.binderInfo
isProp := (← isProp decl.type)
isDecInst := (← forallTelescopeReducing decl.type fun _ type => return type.isAppOf ``Decidable)
}
if decl.binderInfo == .instImplicit then
/- Collect higher order output parameters of this class -/
if let some className ← isClass? decl.type then
if let some outParamPositions := getOutParamPositions? (← getEnv) className then
unless outParamPositions.isEmpty do
let args := decl.type.getAppArgs
for h2 : i in [:args.size] do
if outParamPositions.contains i then
let arg := args[i]
if let some idx := fvars.idxOf? arg then
if (← whnf (← inferType arg)).isForall then
paramInfo := paramInfo.modify idx fun info => { info with higherOrderOutParam := true }
higherOrderOutParams := higherOrderOutParams.insert arg.fvarId!
let resultDeps := collectDeps fvars type
paramInfo := updateHasFwdDeps paramInfo resultDeps
return { resultDeps, paramInfo }
def getFunInfo (fn : Expr) : MetaM FunInfo :=
getFunInfoAux fn none
def getFunInfoNArgs (fn : Expr) (nargs : Nat) : MetaM FunInfo :=
getFunInfoAux fn (some nargs)
def FunInfo.getArity (info : FunInfo) : Nat :=
info.paramInfo.size
end Lean.Meta