Skip to content

Lift tryfinally to macros to fix user code locations #196

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
May 3, 2025
Merged

Lift tryfinally to macros to fix user code locations #196

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 71 additions & 6 deletions src/TimerOutput.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,15 +172,37 @@ Base.@deprecate get_defaultimer get_defaulttimer

# Macro
macro timeit(args...)
return timer_expr(__module__, false, args...)
blocks = timer_expr(__module__, false, args...)
if blocks isa Expr
blocks
else
Expr(:block,
blocks[1], # the timing setup
Expr(:tryfinally,
:($(esc(args[end]))), # the user expr
:($(blocks[2])) # the timing finally
)
)
end
end

macro timeit_debug(args...)
if !isdefined(__module__, :timeit_debug_enabled)
Core.eval(__module__, :(timeit_debug_enabled() = false))
end

return timer_expr(__module__, true, args...)
blocks = timer_expr(__module__, true, args...)
if blocks isa Expr
blocks
else
Expr(:block,
blocks[1], # the timing setup
Expr(:tryfinally,
:($(esc(args[end]))), # the user expr
:($(blocks[2])) # the timing finally
)
)
end
end

function enable_debug_timings(m::Module)
Expand All @@ -204,26 +226,69 @@ function is_func_def(f)
end
end

function _esc(ex)
if isa(ex, Expr)
esc(ex)
else
esc(ex[1]), esc(ex[2])
end
end

function timer_expr(m::Module, is_debug::Bool, ex::Expr)
is_func_def(ex) && return timer_expr_func(m, is_debug, :($(TimerOutputs.DEFAULT_TIMER)), ex)
return esc(_timer_expr(m, is_debug, :($(TimerOutputs).DEFAULT_TIMER), ex))
return _esc(_timer_expr(m, is_debug, :($(TimerOutputs).DEFAULT_TIMER)))
end

function timer_expr(m::Module, is_debug::Bool, label_or_to, ex::Expr)
is_func_def(ex) && return timer_expr_func(m, is_debug, label_or_to, ex)
return esc(_timer_expr(m, is_debug, :($(TimerOutputs).DEFAULT_TIMER), label_or_to, ex))
return _esc(_timer_expr(m, is_debug, :($(TimerOutputs).DEFAULT_TIMER), label_or_to))
end

function timer_expr(m::Module, is_debug::Bool, label::String, ex::Expr)
is_func_def(ex) && return timer_expr_func(m, is_debug, :($(TimerOutputs).DEFAULT_TIMER), ex, label)
return esc(_timer_expr(m, is_debug, :($(TimerOutputs).DEFAULT_TIMER), label, ex))
return _esc(_timer_expr(m, is_debug, :($(TimerOutputs).DEFAULT_TIMER), label))
end

function timer_expr(m::Module, is_debug::Bool, to, label, ex::Expr)
is_func_def(ex) && return timer_expr_func(m, is_debug, to, ex, label)
return esc(_timer_expr(m, is_debug, to, label, ex))
return _esc(_timer_expr(m, is_debug, to, label))
end

# no ex given, so just return before and after for construction in the macro
function _timer_expr(m::Module, is_debug::Bool, to::Union{Symbol, Expr, TimerOutput}, label)
@gensym local_to enabled accumulated_data b₀ t₀ val
timeit_block = quote
$local_to = $to
$enabled = $local_to.enabled
if $enabled
$accumulated_data = $(push!)($local_to, $label)
end
$b₀ = $(gc_bytes)()
$t₀ = $(time_ns)()
end
finally_block = quote
if $enabled
$(do_accumulate!)($accumulated_data, $t₀, $b₀)
$(pop!)($local_to)
end
end

if is_debug
return quote
if $m.timeit_debug_enabled()
$timeit_block
end
end, quote
if $m.timeit_debug_enabled()
$finally_block
end
end
else
return timeit_block, finally_block
end
end

# ex given, so return the whole thing
function _timer_expr(m::Module, is_debug::Bool, to::Union{Symbol, Expr, TimerOutput}, label, ex::Expr)
@gensym local_to enabled accumulated_data b₀ t₀ val
timeit_block = quote
Expand Down