From 9e0ecd2daaba8ac65642671328260401f9079867 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 10 Nov 2024 22:15:39 -0500 Subject: [PATCH 1/3] refactor: use `Meta.isexpr` --- lib/ReactantCore/src/ReactantCore.jl | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/lib/ReactantCore/src/ReactantCore.jl b/lib/ReactantCore/src/ReactantCore.jl index 22ec9f3b9..9b7b34497 100644 --- a/lib/ReactantCore/src/ReactantCore.jl +++ b/lib/ReactantCore/src/ReactantCore.jl @@ -100,12 +100,12 @@ end """ macro trace(expr) expr = macroexpand(__module__, expr) - if expr.head == :(=) - if expr.args[2] isa Expr && expr.args[2].head == :if + if Meta.isexpr(expr, :(=)) + if Meta.isexpr(expr.args[2], :if) return esc(trace_if_with_returns(__module__, expr)) end end - expr.head == :if && return esc(trace_if(__module__, expr)) + Meta.isexpr(expr, :if) && return esc(trace_if(__module__, expr)) return error("Only `if-elseif-else` blocks are currently supported by `@trace`") end @@ -133,7 +133,7 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0) counter = 0 expr = MacroTools.prewalk(expr) do x counter += 1 - if x isa Expr && x.head == :if && counter > 1 + if Meta.isexpr(x, :if) && counter > 1 ex_new, dv, _ = trace_if(mod, x; store_last_line, depth=depth + 1) append!(discard_vars_from_expansion, dv) return ex_new @@ -165,7 +165,7 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0) true_branch_fn_name = gensym(:true_branch) else_block, discard_vars, _ = if length(expr.args) == 3 - if expr.args[3].head != :elseif + if Meta.isexpr(expr.args[3], :elseif) expr.args[3], [], nothing else trace_if(mod, expr.args[3]; store_last_line, depth=depth + 1) @@ -296,10 +296,8 @@ end function error_if_return(expr) return MacroTools.postwalk(expr) do x - if x isa Expr && x.head == :return - error("Cannot use @trace on a block that contains a return statement") - end - return x + Meta.isexpr(x, :return) || return x + error("Cannot use @trace on a block that contains a return statement") end end From 010f7e2dc8af7ff40d0f707b94c6c0c605bb7655 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 10 Nov 2024 22:58:18 -0500 Subject: [PATCH 2/3] feat: support short-circuiting expressions via trace macro --- lib/ReactantCore/src/ReactantCore.jl | 37 ++++++++++++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/lib/ReactantCore/src/ReactantCore.jl b/lib/ReactantCore/src/ReactantCore.jl index 9b7b34497..09035028c 100644 --- a/lib/ReactantCore/src/ReactantCore.jl +++ b/lib/ReactantCore/src/ReactantCore.jl @@ -31,6 +31,7 @@ if no traced value is found inside the expression, then there is no overhead. - `if` conditions (with `elseif` and other niceties) (`@trace if ...`) - `if` statements with a preceeding assignment (`@trace a = if ...`) (note the positioning of the macro needs to be before the assignment and not before the `if`) +- Short circuiting `@trace a && b` or `@trace a || b` ## Special Considerations @@ -100,6 +101,9 @@ end """ macro trace(expr) expr = macroexpand(__module__, expr) + if Meta.isexpr(expr, :(&&), 2) || Meta.isexpr(expr, :(||), 2) + return esc(trace_short_circuit(__module__, expr)) + end if Meta.isexpr(expr, :(=)) if Meta.isexpr(expr.args[2], :if) return esc(trace_if_with_returns(__module__, expr)) @@ -109,6 +113,39 @@ macro trace(expr) return error("Only `if-elseif-else` blocks are currently supported by `@trace`") end +function trace_short_circuit(mod, expr) + if_expr, lhs, varname = generate_if_from_short_circuit(mod, expr) + new_expr = trace_if(mod, if_expr) + return quote + $(varname) = $(lhs) + $(new_expr) + end +end + +function generate_if_from_short_circuit(mod, expr; depth=0) + varname = gensym(:short_circuit_result) + lhs = expr.args[1] + rhs = expr.args[2] + if Meta.isexpr(rhs, :(&&), 2) || Meta.isexpr(rhs, :(||), 2) + rhs = generate_if_from_short_circuit(mod, rhs; depth=depth + 1) + end + if Meta.isexpr(expr, :(&&), 2) + expr = :( + if $(varname) + $(rhs) + end + ) + else + expr = :( + if !$(varname) + $(rhs) + end + ) + end + depth == 0 && return expr, lhs, varname + return :($(varname) = $(lhs); $(expr)) +end + # ... = if ... style expressions function trace_if_with_returns(mod, expr) new_expr, _, all_check_vars = trace_if( From b951b4f4dc3c1e8e4f7fd607acd2f4868dad31a7 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Sun, 10 Nov 2024 23:30:53 -0500 Subject: [PATCH 3/3] fix: incorrect codegen for missing else blocks --- lib/ReactantCore/src/ReactantCore.jl | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/lib/ReactantCore/src/ReactantCore.jl b/lib/ReactantCore/src/ReactantCore.jl index 09035028c..31ac11147 100644 --- a/lib/ReactantCore/src/ReactantCore.jl +++ b/lib/ReactantCore/src/ReactantCore.jl @@ -201,18 +201,24 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0) all_true_branch_vars = true_branch_input_list ∪ true_branch_assignments true_branch_fn_name = gensym(:true_branch) - else_block, discard_vars, _ = if length(expr.args) == 3 + else_block, discard_vars, _, fake_assignments = if length(expr.args) == 3 if Meta.isexpr(expr.args[3], :elseif) - expr.args[3], [], nothing + expr.args[3], [], nothing, :() else - trace_if(mod, expr.args[3]; store_last_line, depth=depth + 1) + (trace_if(mod, expr.args[3]; store_last_line, depth=depth + 1)..., :()) end elseif length(expr.args) == 2 tmp_expr = [] + extra_assignments = [] for var in true_branch_assignments push!(tmp_expr, :($(var) = $(var))) + push!(extra_assignments, :( + if !isdefined($(mod), $(Meta.quot(var))) + $(var) = nothing + end + )) end - Expr(:block, tmp_expr...), [], nothing + Expr(:block, tmp_expr...), [], nothing, Expr(:block, extra_assignments...) else dump(expr) error("This shouldn't happen") @@ -280,6 +286,7 @@ function trace_if(mod, expr; store_last_line=nothing, depth=0) false_branch_fn = :($(false_branch_fn_name) = $(false_branch_fn)) reactant_code_block = quote + $(fake_assignments) $(true_branch_fn) $(false_branch_fn) ($(all_output_vars...),) = $(traced_if)(