90 lines
3.3 KiB
Text
90 lines
3.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
|
||
-/
|
||
prelude
|
||
import Init.Lean.Meta.Basic
|
||
import Init.Lean.Meta.InferType
|
||
|
||
namespace Lean
|
||
namespace Meta
|
||
|
||
@[inline] private def checkFunInfoCache (fn : Expr) (maxArgs? : Option Nat) (k : MetaM FunInfo) : MetaM FunInfo :=
|
||
do s ← get;
|
||
t ← getTransparency;
|
||
match s.cache.funInfo.find ⟨t, fn, maxArgs?⟩ with
|
||
| some finfo => pure finfo
|
||
| none => do
|
||
finfo ← k;
|
||
modify $ fun s => { cache := { funInfo := s.cache.funInfo.insert ⟨t, fn, maxArgs?⟩ finfo, .. s.cache }, .. s };
|
||
pure finfo
|
||
|
||
@[inline] def whenHasVar {α} (e : Expr) (deps : α) (k : α → α) : α :=
|
||
if e.hasFVar then k deps else deps
|
||
|
||
def collectDepsAux (fvars : Array Expr) : Expr → Array Nat → Array Nat
|
||
| e@(Expr.app f a), deps => whenHasVar e deps (collectDepsAux a ∘ collectDepsAux f)
|
||
| e@(Expr.forallE _ _ d b), deps => whenHasVar e deps (collectDepsAux b ∘ collectDepsAux d)
|
||
| e@(Expr.lam _ _ d b), deps => whenHasVar e deps (collectDepsAux b ∘ collectDepsAux d)
|
||
| e@(Expr.letE _ t v b), deps => whenHasVar e deps (collectDepsAux b ∘ collectDepsAux v ∘ collectDepsAux t)
|
||
| Expr.proj _ _ e, deps => collectDepsAux e deps
|
||
| Expr.mdata _ e, deps => collectDepsAux 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
|
||
|
||
def collectDeps (fvars : Array Expr) (e : Expr) : Array Nat :=
|
||
let deps := collectDepsAux fvars 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
|
||
{ hasFwdDeps := true, .. info }
|
||
else
|
||
info
|
||
|
||
@[specialize] def getFunInfoAuxAux
|
||
(whnf : Expr → MetaM Expr)
|
||
(fn : Expr) (maxArgs? : Option Nat) : MetaM FunInfo :=
|
||
checkFunInfoCache fn maxArgs? $ do
|
||
fnType ← inferTypeAux whnf fn;
|
||
forallBoundedTelescope (usingDefault whnf) fnType maxArgs? $ fun fvars type => do
|
||
pinfo ← fvars.size.foldM
|
||
(fun (i : Nat) (pinfo : Array ParamInfo) => do
|
||
let fvar := fvars.get! i;
|
||
decl ← getLocalDecl fvar.fvarId!;
|
||
prop ← isPropAux whnf decl.type;
|
||
let backDeps := collectDeps fvars decl.type;
|
||
let pinfo := updateHasFwdDeps pinfo backDeps;
|
||
pure $ pinfo.push {
|
||
backDeps := backDeps,
|
||
prop := prop,
|
||
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 }
|
||
|
||
@[inline] def getFunInfoAux
|
||
(whnf : Expr → MetaM Expr)
|
||
(fn : Expr) : MetaM FunInfo :=
|
||
getFunInfoAuxAux whnf fn none
|
||
|
||
@[inline] def getFunInfoNArgsAux
|
||
(whnf : Expr → MetaM Expr)
|
||
(fn : Expr) (nargs : Nat) : MetaM FunInfo :=
|
||
getFunInfoAuxAux whnf fn (some nargs)
|
||
|
||
end Meta
|
||
end Lean
|