Fix rewrite of sum(x; init) when x is not a generator (#344)

Author: odow

src/rewrite_generic.jl

@@ -47,6 +47,8 @@ function _is_flatten(expr)
     return Meta.isexpr(expr, :call, 2) && Meta.isexpr(expr.args[2], :flatten)
 end
 
+_is_generator_or_flatten(expr) = _is_generator(expr) || _is_flatten(expr)
+
 function _is_parameters(expr)
     return Meta.isexpr(expr, :call, 3) && Meta.isexpr(expr.args[2], :parameters)
 end
@@ -111,7 +113,7 @@ function _rewrite_generic(stack::Expr, expr::Expr)
     elseif Meta.isexpr(expr.args[2], :(...))
         # If the first argument is a splat.
         return esc(expr), false
-    elseif _is_generator(expr) || _is_flatten(expr) || _is_parameters(expr)
+    elseif _is_generator_or_flatten(expr) || _is_parameters(expr)
         if !(expr.args[1] in (:sum, :Σ, :∑))
             # We don't know what this is. Return the expression and don't let
             # future callers mutate.
@@ -126,7 +128,7 @@ function _rewrite_generic(stack::Expr, expr::Expr)
             # not any of the others.
             p = expr.args[2]
             is_init = length(p.args) == 1 && _is_kwarg(p.args[1], :init)
-            if is_init && expr.args[3] isa Expr
+            if is_init && _is_generator_or_flatten(expr.args[3])
                 # sum(iter ; init) form!
                 # We rewrite only if `iter` is an Expr; if it's just a Symbol,
                 # we don't enter this branch.

test/rewrite_generic.jl

@@ -518,6 +518,52 @@ function test_rewrite_init_symbol()
     return
 end
 
+function test_issue_343()
+    x = [1, 3, 5, 7, 8, 2]
+    y = sum(x)
+    @test MA.@rewrite(sum(x[:]), move_factors_into_sums = false) == y
+    @test MA.@rewrite(sum(x[i] for i in 1:6), move_factors_into_sums = false) ==
+          y
+    @test MA.@rewrite(
+        sum(x[i+j] for i in 1:2:6 for j in 0:1),
+        move_factors_into_sums = false
+    ) == y
+    @test MA.@rewrite(
+        sum(x[i+j] for i in 1:2:6, j in 0:1),
+        move_factors_into_sums = false
+    ) == y
+    # Turn formatting off here so we preserve `init = 0`
+    #!format:off
+    @test MA.@rewrite(sum(x[:], init = 0), move_factors_into_sums = false) == y
+    @test MA.@rewrite(
+        sum(x[i] for i in 1:6, init = 0),
+        move_factors_into_sums = false
+    ) == y
+    @test MA.@rewrite(
+        sum(x[i+j] for i in 1:2:6 for j in 0:1, init = 0),
+        move_factors_into_sums = false
+    ) == y
+    @test MA.@rewrite(
+        sum(x[i+j] for i in 1:2:6, j in 0:1, init = 0),
+        move_factors_into_sums = false
+    ) == y
+    #!format:on
+    @test MA.@rewrite(sum(x[:]; init = 0), move_factors_into_sums = false) == y
+    @test MA.@rewrite(
+        sum(x[i] for i in 1:6; init = 0),
+        move_factors_into_sums = false
+    ) == y
+    @test MA.@rewrite(
+        sum(x[i+j] for i in 1:2:6 for j in 0:1; init = 0),
+        move_factors_into_sums = false
+    ) == y
+    @test MA.@rewrite(
+        sum(x[i+j] for i in 1:2:6, j in 0:1; init = 0),
+        move_factors_into_sums = false
+    ) == y
+    return
+end
+
 end  # module
 
 TestRewriteGeneric.runtests()