78 lines
3 KiB
Text
78 lines
3 KiB
Text
/-
|
||
Copyright (c) 2019 Microsoft Corporation. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Leonardo de Moura
|
||
-/
|
||
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 s ← get
|
||
let t ← getTransparency
|
||
match s.cache.funInfo.find? ⟨t, fn, maxArgs?⟩ with
|
||
| some finfo => pure finfo
|
||
| none => do
|
||
let finfo ← k
|
||
modify fun s => { s with cache := { s.cache with funInfo := s.cache.funInfo.insert ⟨t, fn, maxArgs?⟩ finfo } }
|
||
pure 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 : Expr → Array Nat → Array Nat
|
||
| e@(Expr.app f a _), deps => whenHasVar e deps (visit a ∘ visit f)
|
||
| e@(Expr.forallE _ d b _), deps => whenHasVar e deps (visit b ∘ visit d)
|
||
| e@(Expr.lam _ d b _), deps => whenHasVar e deps (visit b ∘ visit d)
|
||
| e@(Expr.letE _ t v b _), deps => whenHasVar e deps (visit b ∘ visit v ∘ visit t)
|
||
| Expr.proj _ _ e _, deps => visit e deps
|
||
| Expr.mdata _ e _, deps => visit e deps
|
||
| e@(Expr.fvar _ _), deps =>
|
||
match fvars.indexOf? e with
|
||
| none => deps
|
||
| some i => if deps.contains i.val then deps else deps.push i.val
|
||
| _, deps => 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
|
||
withTransparency TransparencyMode.default $
|
||
forallBoundedTelescope fnType maxArgs? fun fvars type => do
|
||
let mut pinfo := #[]
|
||
for i in [:fvars.size] do
|
||
let fvar := fvars[i]
|
||
let decl ← getFVarLocalDecl fvar
|
||
let backDeps := collectDeps fvars decl.type
|
||
let pinfo := updateHasFwdDeps pinfo backDeps
|
||
pinfo := pinfo.push {
|
||
backDeps := backDeps,
|
||
implicit := decl.binderInfo == BinderInfo.implicit,
|
||
instImplicit := decl.binderInfo == BinderInfo.instImplicit }
|
||
let resultDeps := collectDeps fvars type
|
||
let pinfo := updateHasFwdDeps pinfo resultDeps
|
||
pure { resultDeps := resultDeps, paramInfo := pinfo }
|
||
|
||
def getFunInfo (fn : Expr) : MetaM FunInfo :=
|
||
getFunInfoAux fn none
|
||
|
||
def getFunInfoNArgs (fn : Expr) (nargs : Nat) : MetaM FunInfo :=
|
||
getFunInfoAux fn (some nargs)
|
||
|
||
end Lean.Meta
|