-
Notifications
You must be signed in to change notification settings - Fork 1
/
LocalScopeBenchmarks.jl
70 lines (59 loc) · 1.73 KB
/
LocalScopeBenchmarks.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
module LocalScopeBenchmarks
import BenchmarkTools
using MacroTools: MacroTools, postwalk, @capture
using OrderedCollections: OrderedDict
function collect_symbols(expr)
assignments = OrderedDict{Symbol, Expr}()
postwalk(expr) do x
if x isa Symbol
assignments[x] = Expr(:$, x)
end
end
assignments
end
function parse_setup(setup::Expr)
assignments = OrderedDict()
postwalk(setup) do x
if @capture(x, _a = _b)
assignments[a] = b
end
end
assignments
end
function lower_setup(assignments::AbstractDict)
Expr(:block, [Expr(:(=), k, v) for (k, v) in assignments]...)
end
function parse_params(kwargs)
params_dict = OrderedDict((@assert x.head == :kw; x.args[1] => x.args[2]) for x in kwargs)
end
function lower_params(params::AbstractDict)
[Expr(:kw, k, v) for (k, v) in params]
end
function interpolate_locals_into_setup(args...)
core, kwargs = BenchmarkTools.prunekwargs(args...)
params = parse_params(kwargs)
setup_assignments = parse_setup(get(() -> Expr(:block), params, :setup))
local_assignments = collect_symbols(core)
setup = merge(local_assignments, setup_assignments)
params[:setup] = lower_setup(setup)
core, lower_params(params)
end
macro benchmark(args...)
core, params = interpolate_locals_into_setup(args...)
quote
BenchmarkTools.@benchmark($(core), $(params...))
end
end
macro btime(expr, args...)
core, params = interpolate_locals_into_setup(args...)
quote
BenchmarkTools.@btime($(core), $(params...))
end
end
macro belapsed(expr, args...)
core, params = interpolate_locals_into_setup(args...)
quote
BenchmarkTools.@belapsed($(core), $(params...))
end
end
end