Fix rewrite of +(x, *(y...)) into add_mul (#225)

* Fix rewrite of +(x, *(y...)) into add_mul

* Add test

* Add more tests

src/rewrite_generic.jl

@@ -121,6 +121,21 @@ function _rewrite_generic(stack::Expr, expr::Expr)
         @assert length(expr.args) > 1
         if length(expr.args) == 2  # +(arg)
             return _rewrite_generic(stack, expr.args[2])
+        elseif length(expr.args) == 3 && _is_call(expr.args[3], :*)
+            # +(x, *(y...)) => add_mul(x, y...)
+            x, is_mutable = _rewrite_generic(stack, expr.args[2])
+            rhs = if is_mutable
+                Expr(:call, operate!!, add_mul, x)
+            else
+                Expr(:call, operate, add_mul, x)
+            end
+            for i in 2:length(expr.args[3].args)
+                yi, _ = _rewrite_generic(stack, expr.args[3].args[i])
+                push!(rhs.args, yi)
+            end
+            root = gensym()
+            push!(stack.args, :($root = $rhs))
+            return root, true
         end
         return _rewrite_generic_to_nested_op(stack, expr, add_mul)
     elseif expr.args[1] == :-

test/rewrite_generic.jl

@@ -324,6 +324,26 @@ function test_rewrite_kw_in_ref()
     return
 end
 
+function test_rewrite_expression()
+    x = [1.2]
+    @test MA.@rewrite(x + 2 * x, move_factors_into_sums = false) == 3x
+    @test MA.@rewrite(x + *(2, x, 3), move_factors_into_sums = false) ==
+          x + *(2, x, 3)
+    y = 1.2
+    @test MA.@rewrite(
+        sum(y for i in 1:2) + 2y,
+        move_factors_into_sums = false
+    ) == 4y
+    @test MA.@rewrite(
+        sum(y for i in 1:2) + y * y,
+        move_factors_into_sums = false
+    ) == 2y + y^2
+    @test MA.@rewrite(y + 2 * y, move_factors_into_sums = false) == 3y
+    @test MA.@rewrite(y + *(2, y, 3), move_factors_into_sums = false) ==
+          y + *(2, y, 3)
+    return
+end
+
 end  # module
 
 TestRewriteGeneric.runtests()