feat: add ForLoop

Author: AayushSabharwal

src/code.jl

@@ -5,7 +5,7 @@ using StaticArrays, SparseArrays, LinearAlgebra, NaNMath, SpecialFunctions,
 
 export toexpr, Assignment, (←), Let, Func, DestructuredArgs, LiteralExpr,
        SetArray, MakeArray, MakeSparseArray, MakeTuple, AtIndex,
-       SpawnFetch, Multithreaded, cse
+       SpawnFetch, Multithreaded, ForLoop, cse
 
 import ..SymbolicUtils
 import ..SymbolicUtils.Rewriters
@@ -705,6 +705,27 @@ function toexpr(exp::LiteralExpr, st)
     recurse_expr(exp.ex, st)
 end
 
+"""
+    ForLoop(itervar, range, body)
+
+Generate a `for` loop of the form
+```julia
+for itervar in range
+    body
+end
+```
+"""
+struct ForLoop <: CodegenPrimitive
+    itervar
+    range
+    body
+end
+
+function toexpr(f::ForLoop, st)
+    :(for $(toexpr(f.itervar, st)) in $(toexpr(f.range, st))
+        $(toexpr(f.body, st))
+    end)
+end
 
 ### Code-related utilities
 
@@ -935,4 +956,10 @@ function cse!(x::SpawnFetch{T}, state::CSEState) where {T}
     return SpawnFetch{T}(map(cse, x.exprs), cse!(x.args, state), x.combine)
 end
 
+function cse!(x::ForLoop, state::CSEState)
+    # cse the range with current scope, CSE the body with a new scope
+    new_state = new_scope(state)
+    return ForLoop(x.itervar, cse!(x.range, state), apply_cse(cse!(x.body, new_state), new_state))
+end
+
 end

test/code.jl

@@ -305,3 +305,21 @@ end
         @test arr[i] == 0
     end
 end
+
+@testset "`ForLoop`" begin
+    @syms a b c::Array
+    ex = ForLoop(a, term(range, b^2, b^2 + 3), SetArray(false, c, [AtIndex(a, a + 1)]))
+    expr = quote
+        let b = 2, c = zeros(Int, 10)
+            $(toexpr(ex))
+            c
+        end
+    end
+    arr = eval(expr)
+    @test arr[4] == 5
+    @test arr[5] == 6
+    @test arr[6] == 7
+    @test arr[7] == 8
+    @test all(iszero, arr[1:3])
+    @test all(iszero, arr[8:end])
+end

test/cse.jl

@@ -201,3 +201,24 @@ end
     rgf = @RuntimeGeneratedFunction(toexpr(innerfn))
     @test rgf(5, sin(1)) == trueval
 end
+
+@testset "ForLoop" begin
+    @syms a b c::Array
+    ex = ForLoop(a, term(range, b^2, b^2 + 3), SetArray(false, c, [AtIndex(a, a^2 + sin(a^2))]))
+    csex = cse(ex)
+    @test findfirst(isequal(csex.body.range), Code.lhs.(csex.pairs)) !== nothing
+    @test findfirst(isequal(csex.body.body.body.elems[1].elem), Code.lhs.(csex.body.body.pairs)) !== nothing
+    expr = quote
+        let b = 2, c = zeros(10)
+            $(toexpr(ex))
+            c
+        end
+    end
+    arr = eval(expr)
+    @test arr[4] == 4^2 + sin(4^2)
+    @test arr[5] == 5^2 + sin(5^2)
+    @test arr[6] == 6^2 + sin(6^2)
+    @test arr[7] == 7^2 + sin(7^2)
+    @test all(iszero, arr[1:3])
+    @test all(iszero, arr[8:end])
+end