try lifting tryfinally to macros

Author: IanButterworth

src/TimerOutput.jl

@@ -172,7 +172,18 @@ 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...)
@@ -204,24 +215,32 @@ 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), ex))
 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, ex))
 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, ex))
 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, ex))
 end
 
 function _timer_expr(m::Module, is_debug::Bool, to::Union{Symbol, Expr, TimerOutput}, label, ex::Expr)
@@ -234,27 +253,26 @@ function _timer_expr(m::Module, is_debug::Bool, to::Union{Symbol, Expr, TimerOut
         end
         $b₀ = $(gc_bytes)()
         $t₀ = $(time_ns)()
-        $(Expr(:tryfinally,
-            :($val = $ex),
-            quote
-                if $enabled
-                    $(do_accumulate!)($accumulated_data, $t₀, $b₀)
-                    $(pop!)($local_to)
-                end
-            end))
-        $val
+    end
+    finally_block = quote
+        if $enabled
+            $(do_accumulate!)($accumulated_data, $t₀, $b₀)
+            $(pop!)($local_to)
+        end
     end
 
     if is_debug
-        return Base.remove_linenums!(quote
+        return quote
             if $m.timeit_debug_enabled()
                 $timeit_block
-            else
-                $ex
             end
-        end)
+        end, quote
+            if $m.timeit_debug_enabled()
+                $finally_block
+            end
+        end
     else
-        return Base.remove_linenums!(timeit_block)
+        return timeit_block, finally_block
     end
 end