vibe code a macro that annotates all lines in a section

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)

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₀

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!
 
 

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()