@trace if return nothing in branches (#1311)

* return nothing if the value if the result value in an if-statement is not required.

* bump

* don't create missing traced value if the variable exists before the if statement already.

* actually fix

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* test

* add test

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Paul Berg <[email protected]>

Author: 3 people

lib/ReactantCore/src/ReactantCore.jl

@@ -414,7 +414,10 @@ function trace_if(expr; store_last_line=nothing, depth=0, track_numbers)
             $(store_last_line) = $(true_last_line)
         end
     else
-        expr.args[2]
+        quote
+            $(expr.args[2])
+            nothing # explicitly return nothing to prevent branches from returning different types
+        end
     end
 
     true_branch_symbols = ExpressionExplorer.compute_symbols_state(true_block)
@@ -457,7 +460,10 @@ function trace_if(expr; store_last_line=nothing, depth=0, track_numbers)
             $(store_last_line) = $(false_last_line)
         end
     else
-        else_block
+        quote
+            $else_block
+            nothing # explicitly return nothing to prevent branches from returning different types
+        end
     end
 
     false_branch_symbols = ExpressionExplorer.compute_symbols_state(false_block)
@@ -473,10 +479,12 @@ function trace_if(expr; store_last_line=nothing, depth=0, track_numbers)
 
     all_vars = all_input_vars ∪ all_output_vars
 
-    non_existant_true_branch_vars = setdiff(all_output_vars, all_true_branch_vars)
+    non_existent_true_branch_vars = setdiff(
+        all_output_vars, all_true_branch_vars, all_input_vars
+    )
     true_branch_extras = Expr(
         :block,
-        [:($(var) = $(MissingTracedValue)()) for var in non_existant_true_branch_vars]...,
+        [:($(var) = $(MissingTracedValue)()) for var in non_existent_true_branch_vars]...,
     )
 
     true_branch_fn = :(($(all_input_vars...),) -> begin
@@ -489,12 +497,12 @@ function trace_if(expr; store_last_line=nothing, depth=0, track_numbers)
     )
     true_branch_fn = :($(true_branch_fn_name) = $(true_branch_fn))
 
-    non_existant_false_branch_vars = setdiff(
-        setdiff(all_output_vars, all_false_branch_vars), all_input_vars
+    non_existent_false_branch_vars = setdiff(
+        all_output_vars, all_false_branch_vars, all_input_vars
     )
     false_branch_extras = Expr(
         :block,
-        [:($(var) = $(MissingTracedValue)()) for var in non_existant_false_branch_vars]...,
+        [:($(var) = $(MissingTracedValue)()) for var in non_existent_false_branch_vars]...,
     )
 
     false_branch_fn = :(($(all_input_vars...),) -> begin

test/control_flow.jl

@@ -934,3 +934,36 @@ end
 
     @test h_ra ≈ h
 end
+
+function different_branch_returns(cond, a, b)
+    @trace if cond
+        a .= sin.(a)
+        nothing
+    else
+        b .= sin.(b)
+        nothing
+    end
+    return a, b
+end
+
+@testset "one branch mutates variable" begin
+    cond = true
+    a = 3 .* ones(Float32, 2, 3)
+    b = 4 .* ones(Float64, 2, 3)
+
+    cond_ra = ConcreteRNumber{Bool}(cond)
+    a_ra = Reactant.to_rarray(a)
+    b_ra = Reactant.to_rarray(b)
+
+    result_ra = @jit(different_branch_returns(cond_ra, a_ra, b_ra))
+    result = different_branch_returns(cond, a, b)
+
+    @test result_ra[1] == a_ra
+    @test result_ra[2] == b_ra
+
+    @test result[1] == a
+    @test result[2] == b
+
+    @test a_ra ≈ a
+    @test b_ra ≈ b
+end