feat: allow SetArray to return the array being set

Author: AayushSabharwal

src/code.jl

@@ -416,10 +416,11 @@ end
     inbounds::Bool
     arr
     elems  # Either iterator of Pairs or just an iterator
+    return_arr::Bool
 end
 
 """
-    SetArray(inbounds, arr, elems)
+    SetArray(inbounds::Bool, arr, elems[, return_arr::Bool])
 
 An expression representing setting of elements of `arr`.
 
@@ -430,9 +431,14 @@ is performed in its place.
 
 `inbounds` is a boolean flag, `true` surrounds the resulting expression
 in an `@inbounds`.
+
+`return_arr` is a flag which controls whether the generated `begin..end` block
+returns the `arr`. Defaults to `false`, in which case the block returns `nothing`.
 """
 SetArray
 
+SetArray(inbounds, arr, elems) = SetArray(inbounds, arr, elems, false)
+
 @matchable struct AtIndex <: CodegenPrimitive
     i
     elem
@@ -446,7 +452,7 @@ function toexpr(s::SetArray, st)
     ex = quote
         $([:($(toexpr(s.arr, st))[$(ex isa AtIndex ? toexpr(ex.i, st) : i)] = $(toexpr(ex, st)))
            for (i, ex) in enumerate(s.elems)]...)
-        nothing
+        $(s.return_arr ? toexpr(s.arr, st) : nothing)
     end
     s.inbounds ? :(@inbounds $ex) : ex
 end
@@ -906,7 +912,7 @@ function cse!(x::MakeArray, state::CSEState)
 end
 
 function cse!(x::SetArray, state::CSEState)
-    return SetArray(x.inbounds, x.arr, cse!(x.elems, state))
+    return SetArray(x.inbounds, x.arr, cse!(x.elems, state), x.return_arr)
 end
 
 function cse!(x::MakeSparseArray, state::CSEState)

test/code.jl

@@ -323,3 +323,21 @@ end
     @test all(iszero, arr[1:3])
     @test all(iszero, arr[8:end])
 end
+
+@testset "`SetArray` with `return_arr`" begin
+    @syms a b c::Array
+    ex = SetArray(false, c, [3, 2, 1], false)
+    expr = quote
+        let b = 2, c = zeros(Int, 3)
+            $(toexpr(ex))
+        end
+    end
+    @test eval(expr) === nothing
+    ex = SetArray(false, c, [3, 2, 1], true)
+    expr = quote
+        let b = 2, c = zeros(Int, 3)
+            $(toexpr(ex))
+        end
+    end
+    @test eval(expr) == [3, 2, 1]
+end

test/cse.jl

@@ -107,7 +107,7 @@ end
     marr = MakeArray(arr, Array)
     sparr = sparse([1, 2, 3, 4], [1, 2, 3, 4], vec(arr))
     msparr = MakeSparseArray(sparr)
-    sarr = SetArray(false, :buffer, [[a^2 + c^2], AtIndex(3, arr), AtIndex(4, msparr)])
+    sarr = SetArray(false, :buffer, [[a^2 + c^2], AtIndex(3, arr), AtIndex(4, msparr)], true)
 
     csex = cse(sarr)
     # test that simple array is CSEd
@@ -127,7 +127,6 @@ end
     expr = quote
         let a = 1, b = 2, c = 3, buffer = Any[0, "A", 0, 0]
             $(toexpr(csex))
-            buffer
         end
     end
     val = eval(expr)