respect scopping rules in for

lib/ReactantCore/src/ReactantCore.jl

@@ -156,8 +156,15 @@ function trace_for(mod, expr)
         :tuple, :(Reactant.promote_to(Reactant.TracedRNumber{Int}, 0)), external_syms...
     )
 
+    cond_val(s) = :(isdefined($(mod), $(QuoteNode(s))) ? $s : nothing)
+
+    locals = Expr[
+        [Expr(:(=), s, cond_val(s)) for s in external_syms]...,
+        :(args = $(args_init)),
+    ]
+
     reactant_code_block = quote
-        let args = $(args_init)
+        let $(locals...)
             cond_fn =
                 $(all_syms) -> begin
                     local num_iters = div($limit - $start, $step, RoundDown)
@@ -180,7 +187,7 @@ function trace_for(mod, expr)
     end
 
     return quote
-        if any($(is_traced), $(Expr(:tuple, all_syms.args[(begin + 1):end]...)))
+        if any($(is_traced), $(Expr(:tuple, cond_val.(all_syms.args[(begin + 1):end])...)))
             $(reactant_code_block)
         else
             $(expr)

src/ControlFlow.jl

@@ -118,16 +118,21 @@ function ReactantCore.traced_while(
 
     result_0 = in_tys
 
-    operands = MLIR.IR.Value[v.mlir_data for v in traced_args]
+    operands = MLIR.IR.Value[v.mlir_data for v in traced_args if v isa TracedType]
 
     while_compiled = MLIR.Dialects.stablehlo.while_(
         operands; result_0, cond=cond_reg, body=body_reg
     )
 
-    return map(enumerate(traced_args)) do (i, res)
-        res.mlir_data = MLIR.IR.result(while_compiled, i)
-        return res
+    residx = 1
+    for res in traced_args
+        if res isa TracedType
+            res.mlir_data = MLIR.IR.result(while_compiled, residx)
+            residx += 1
+        end
     end
+
+    return traced_args
 end
 
 function take_region(compiled_fn)

test/control_flow.jl

@@ -565,3 +565,18 @@ end
 
     @test @jit(for_ref_outer(x_ra)) ≈ for_ref_outer(x)
 end
+
+function for_inner_scope(x)
+    @trace for i in 1:10
+        s = sum(x)
+        x = x / s
+    end
+    return x
+end
+
+@testset "for: inner scope" begin
+    x = randn(Float64, 10)
+    x_ra = Reactant.to_rarray(x)
+
+    @test @jit(for_inner_scope(x_ra)) ≈ for_inner_scope(x)
+end