61 lines
1.8 KiB
Text
61 lines
1.8 KiB
Text
/-
|
||
Copyright (c) 2024 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 import Lean.Util.PtrSet
|
||
|
||
public section
|
||
|
||
namespace Lean.Expr
|
||
namespace NumApps
|
||
|
||
unsafe structure State where
|
||
visited : PtrSet Expr := mkPtrSet
|
||
counters : NameMap Nat := {}
|
||
|
||
unsafe abbrev M := StateM State
|
||
|
||
unsafe def visit (e : Expr) : M Unit :=
|
||
unless (← get).visited.contains e do
|
||
modify fun s => { s with visited := s.visited.insert e }
|
||
match e with
|
||
| .forallE _ d b _ => visit d; visit b
|
||
| .lam _ d b _ => visit d; visit b
|
||
| .mdata _ b => visit b
|
||
| .letE _ t v b _ => visit t; visit v; visit b
|
||
| .app .. => e.withApp fun f args => do
|
||
if let .const declName _ := f then
|
||
let c := (← get).counters.find? declName |>.getD 0
|
||
modify fun s => { s with counters := s.counters.insert declName (c+1) }
|
||
visit f
|
||
args.forM visit
|
||
| .proj _ _ b => visit b
|
||
| _ => return ()
|
||
|
||
unsafe def main (e : Expr) : NameMap Nat :=
|
||
let (_, s) := NumApps.visit e |>.run {}
|
||
s.counters
|
||
|
||
end NumApps
|
||
|
||
/--
|
||
Returns the number of applications for each declaration used in `e`.
|
||
|
||
This operation is performed in `IO` because the result depends on the memory representation of the object.
|
||
|
||
Note: Use this function primarily for diagnosing performance issues.
|
||
-/
|
||
def numApps (e : Expr) (threshold : Nat := 0) : IO (Array (Name × Nat)) := do
|
||
let counters := unsafe NumApps.main e
|
||
let mut result := #[]
|
||
for (declName, num) in counters do
|
||
if num > threshold then
|
||
result := result.push (declName, num)
|
||
return result.qsort fun a b => a.2 > b.2
|
||
|
||
end Lean.Expr
|