From c463ec58bdac20310bfc24015d956881ffea2f05 Mon Sep 17 00:00:00 2001 From: Kristoffer Carlsson Date: Thu, 22 Jan 2026 13:02:10 +0100 Subject: [PATCH] vibe code a macro that annotates all lines in a section --- README.md | 11 +++ src/TimerOutput.jl | 161 ++++++++++++++++++++++++++++++++++++++++++++ src/TimerOutputs.jl | 2 +- test/runtests.jl | 19 ++++++ 4 files changed, 192 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 586f030..99d5714 100644 --- a/README.md +++ b/README.md @@ -105,6 +105,17 @@ end funcdef(2) +# Annotate each line in a block or function definition +@timeit_all to function line_profile(n) + x = 0 + for i in 1:n + x += i + end + x +end + +line_profile(10) + # Or to instrument an existing function: foo(x) = x + 1 timed_foo = to(foo) diff --git a/src/TimerOutput.jl b/src/TimerOutput.jl index 99f5244..556c1ab 100644 --- a/src/TimerOutput.jl +++ b/src/TimerOutput.jl @@ -218,6 +218,10 @@ macro timeit_debug(args...) end end +macro timeit_all(args...) + return timeit_all_expr(__source__, __module__, args...) +end + function enable_debug_timings(m::Module) return if !getfield(m, :timeit_debug_enabled)() Core.eval(m, :(timeit_debug_enabled() = true)) @@ -230,6 +234,7 @@ function disable_debug_timings(m::Module) end timer_expr(args...) = throw(ArgumentError("invalid macro usage for @timeit, use as @timeit [to] label codeblock")) +timeit_all_expr(args...) = throw(ArgumentError("invalid macro usage for @timeit_all, use as @timeit_all [to] [label] codeblock")) function is_func_def(f) if isa(f, Expr) && (f.head === :function || Base.is_short_function_def(f)) @@ -267,6 +272,26 @@ function timer_expr(source::LineNumberNode, m::Module, is_debug::Bool, to, label return _esc(_timer_expr(source, m, is_debug, to, label)) end +function timeit_all_expr(source::LineNumberNode, m::Module, ex::Expr) + is_func_def(ex) && return timeit_all_expr_func(source, m, :($(TimerOutputs).DEFAULT_TIMER), ex) + return _esc(_timeit_all_expr(source, m, :($(TimerOutputs).DEFAULT_TIMER), nothing, ex)) +end + +function timeit_all_expr(source::LineNumberNode, m::Module, label_or_to, ex::Expr) + is_func_def(ex) && return timeit_all_expr_func(source, m, label_or_to, ex) + return _esc(_timeit_all_expr(source, m, label_or_to, nothing, ex)) +end + +function timeit_all_expr(source::LineNumberNode, m::Module, label::String, ex::Expr) + is_func_def(ex) && return timeit_all_expr_func(source, m, :($(TimerOutputs).DEFAULT_TIMER), ex, label) + return _esc(_timeit_all_expr(source, m, :($(TimerOutputs).DEFAULT_TIMER), label, ex)) +end + +function timeit_all_expr(source::LineNumberNode, m::Module, to, label, ex::Expr) + is_func_def(ex) && return timeit_all_expr_func(source, m, to, ex, label) + return _esc(_timeit_all_expr(source, m, to, label, ex)) +end + # no ex given, so just return before and after for construction in the macro function _timer_expr(source::LineNumberNode, m::Module, is_debug::Bool, to::Union{Symbol, Expr, TimerOutput}, label) @gensym local_to enabled accumulated_data b₀ t₀ val @@ -366,6 +391,142 @@ function timer_expr_func(source::LineNumberNode, m::Module, is_debug::Bool, to, return esc(combinedef(def)) end +function timeit_all_expr_func(source::LineNumberNode, m::Module, to, expr::Expr, label = nothing) + expr = macroexpand(m, expr) + def = splitdef(expr) + label === nothing && (label = line_label(source, expr)) + def[:body] = _timer_expr(source, m, false, to, label, annotate_lines(def[:body], source, m, to)) + return esc(combinedef(def)) +end + +function _timeit_all_expr(source::LineNumberNode, m::Module, to, label, ex::Expr) + body = annotate_lines(ex, source, m, to) + return label === nothing ? body : _timer_expr(source, m, false, to, label, body) +end + +function annotate_lines(ex, source::LineNumberNode, m::Module, to) + return ex isa Expr && ex.head === :block ? annotate_block(ex, source, m, to) : instrument_statement(ex, source, m, to) +end + +function annotate_block(ex::Expr, source::LineNumberNode, m::Module, to) + args = Any[] + current_source = source + for arg in ex.args + if arg isa LineNumberNode + current_source = arg + elseif arg isa Expr && arg.head === :line && !isempty(arg.args) + line = arg.args[1] + file = length(arg.args) > 1 ? arg.args[2] : current_source.file + current_source = LineNumberNode(line, file) + elseif arg isa Expr && arg.head === :meta + push!(args, arg) + else + push!(args, instrument_statement(arg, current_source, m, to)) + end + end + return Expr(:block, args...) +end + +function instrument_statement(stmt, source::LineNumberNode, m::Module, to) + if !(stmt isa Expr) + expr = Expr(:block, stmt) + return _timer_expr(source, m, false, to, line_label(source, stmt), expr) + end + if is_func_def(stmt) + return timeit_all_expr_func(source, m, to, stmt) + end + if stmt.head === :meta + return stmt + end + if is_macrocall_named(stmt, Symbol("@notimeit")) + return stmt + end + stmt = annotate_subexprs(stmt, source, m, to) + return _timer_expr(source, m, false, to, line_label(source, stmt), stmt) +end + +function annotate_subexprs(ex::Expr, source::LineNumberNode, m::Module, to) + head = ex.head + if head === :block + return annotate_block(ex, source, m, to) + elseif head === :for || head === :while + return Expr(head, ex.args[1], annotate_lines(ex.args[2], source, m, to)) + elseif head === :if + then_block = annotate_lines(ex.args[2], source, m, to) + if length(ex.args) >= 3 + else_arg = ex.args[3] + else_block = else_arg isa Expr && else_arg.head === :elseif ? + annotate_subexprs(else_arg, source, m, to) : + annotate_lines(else_arg, source, m, to) + return Expr(:if, ex.args[1], then_block, else_block) + end + return Expr(:if, ex.args[1], then_block) + elseif head === :elseif + cond_block = ex.args[1] + then_block = annotate_lines(ex.args[2], source, m, to) + else_arg = ex.args[3] + else_block = else_arg isa Expr && else_arg.head === :elseif ? + annotate_subexprs(else_arg, source, m, to) : + annotate_lines(else_arg, source, m, to) + return Expr(:elseif, cond_block, then_block, else_block) + elseif head === :let + if isempty(ex.args) + return ex + end + body = annotate_lines(ex.args[end], source, m, to) + return Expr(:let, ex.args[1:end-1]..., body) + elseif head === :try + return annotate_try(ex, source, m, to) + elseif head === :quote + return ex + else + return ex + end +end + +function annotate_try(ex::Expr, source::LineNumberNode, m::Module, to) + args = Any[] + for arg in ex.args + if arg isa Expr && arg.head === :catch + body = annotate_lines(arg.args[end], source, m, to) + push!(args, Expr(:catch, arg.args[1:end-1]..., body)) + elseif arg isa Expr && arg.head === :finally + body = annotate_lines(arg.args[end], source, m, to) + push!(args, Expr(:finally, arg.args[1:end-1]..., body)) + elseif arg isa Expr + push!(args, annotate_lines(arg, source, m, to)) + else + push!(args, arg) + end + end + return Expr(:try, args...) +end + +function is_macrocall_named(ex::Expr, name::Symbol) + if ex.head !== :macrocall || isempty(ex.args) + return false + end + head = ex.args[1] + return head === name || (head isa Expr && head.head === :. && head.args[end] === name) +end + +function line_label(source::LineNumberNode, ex) + file = source.file === nothing ? "unknown" : string(source.file) + file = Base.basename(file) + return string(file, ":", source.line, " ", expr_summary(ex)) +end + +function expr_summary(ex) + if ex isa Expr + ex = Base.remove_linenums!(copy(ex)) + s = sprint(io -> Base.show_unquoted(io, ex, 0, 0)) + else + s = string(ex) + end + first_line = split(s, '\n', limit = 2)[1] + return strip(first_line) +end + function do_accumulate!(accumulated_data, t₀, b₀) accumulated_data.time += time_ns() - t₀ accumulated_data.allocs += gc_bytes() - b₀ diff --git a/src/TimerOutputs.jl b/src/TimerOutputs.jl index 759677a..10307b1 100644 --- a/src/TimerOutputs.jl +++ b/src/TimerOutputs.jl @@ -4,7 +4,7 @@ using ExprTools import Base: show, time_ns export TimerOutput, @timeit, @timeit_debug, reset_timer!, print_timer, timeit, - enable_timer!, disable_timer!, @notimeit, get_timer, + enable_timer!, disable_timer!, @notimeit, get_timer, @timeit_all, begin_timed_section!, end_timed_section! diff --git a/test/runtests.jl b/test/runtests.jl index 1cfb857..db5f861 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -273,6 +273,25 @@ end @test ncalls(DEFAULT_TIMER["baz"]) == 2 @test ncalls(DEFAULT_TIMER["quz"]) == 2 + to_lines = TimerOutput() + @timeit_all to_lines function timeit_all_test(n) + x = 0 + for i in 1:n + x += i + end + x + end + @test timeit_all_test(3) == 6 + filename = basename(@__FILE__) + func_keys = [k for k in keys(to_lines.inner_timers) + if occursin(filename, k) && occursin("function timeit_all_test", k)] + @test length(func_keys) == 1 + func_timer = to_lines.inner_timers[func_keys[1]] + for_keys = [k for k in keys(func_timer.inner_timers) if occursin("for i = 1:n", k)] + @test length(for_keys) == 1 + for_timer = func_timer.inner_timers[for_keys[1]] + @test any(occursin("x += i", k) for k in keys(for_timer.inner_timers)) + function foo() reset_timer!() @timeit "asdf" bar()