Fix unary minus operators in rewrite

src/rewrite_generic.jl

@@ -191,7 +191,15 @@ function _rewrite_generic(stack::Expr, expr::Expr)
         # -(args...) => sub_mul(sub_mul(arg1, arg2), arg3)
         @assert length(expr.args) > 1
         if length(expr.args) == 2  # -(arg)
-            return _rewrite_generic(stack, Expr(:call, :*, -1, expr.args[2]))
+            value, is_mutable = _rewrite_generic(stack, expr.args[2])
+            root = gensym()
+            rhs = if is_mutable
+                Expr(:call, operate!!, *, value, -1)
+            else
+                Expr(:call, operate, *, -1, value)
+            end
+            push!(stack.args, :($root = $rhs))
+            return root, true
         end
         return _rewrite_generic_to_nested_op(stack, expr, sub_mul)
     elseif expr.args[1] == :*
@@ -278,6 +286,7 @@ function _rewrite_generic_to_nested_op(stack, expr, op; broadcast::Bool = false)
 end
 
 _is_call(expr, op) = Meta.isexpr(expr, :call) && expr.args[1] == op
+_is_call(expr, op, n) = Meta.isexpr(expr, :call, 1 + n) && expr.args[1] == op
 
 """
     _rewrite_generic_generator(stack::Expr, op::Symbol, expr::Expr)
@@ -322,6 +331,12 @@ function _rewrite_generic_generator(
             push!(rhs.args, value)
         end
         push!(new_stack.args, :($root = $rhs))
+    elseif op == :+ && _is_call(expr.args[1], :-, 1)
+        # Optimization time! Instead of operate!!(+, root, -(arg)), rewrite this
+        # as operate!!(sub_mul, root, arg)
+        value, _ = _rewrite_generic(new_stack, expr.args[1].args[2])
+        rhs = Expr(:call, operate!!, sub_mul, root, value)
+        push!(new_stack.args, :($root = $rhs))
     elseif is_flatten
         # The first argument is itself a generator
         _rewrite_generic_generator(new_stack, op, expr.args[1], root)

test/rewrite_generic.jl

@@ -465,6 +465,39 @@ function test_return_is_mutable()
     return
 end
 
+function test_rewrite_sum_unary_minus()
+    x = big.(1:3)
+    y = MA.@rewrite(sum(-x[i] for i in 1:3), move_factors_into_sums = false)
+    @test y == big(-6)
+    return
+end
+
+function test_allocations_rewrite_unary_minus()
+    N = 100
+    x = big.(1:N)
+    a = big(-1)
+    # precompilaton
+    MA.@rewrite(sum(-x[i] for i in 1:N), move_factors_into_sums = false)
+    MA.@rewrite(-sum(x[i] for i in 1:N), move_factors_into_sums = false)
+    MA.@rewrite(sum(a * x[i] for i in 1:N), move_factors_into_sums = false)
+    sum(-x[i] for i in 1:N)
+    total = @allocated sum(-x[i] for i in 1:N)
+    # actual
+    value = @allocated(
+        MA.@rewrite(sum(-x[i] for i in 1:N), move_factors_into_sums = false),
+    )
+    @test value < total
+    value = @allocated(
+        MA.@rewrite(-sum(x[i] for i in 1:N), move_factors_into_sums = false),
+    )
+    @test value < total
+    value = @allocated(
+        MA.@rewrite(sum(a * x[i] for i in 1:N), move_factors_into_sums = false),
+    )
+    @test value < total
+    return
+end
+
 end  # module
 
 TestRewriteGeneric.runtests()