Add test

jump-dev · Sep 1, 2023 · 115d97c · 115d97c
1 parent 72dbd20
commit 115d97c
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git a/src/rewrite_generic.jl b/src/rewrite_generic.jl
@@ -124,14 +124,18 @@ function _rewrite_generic(stack::Expr, expr::Expr)
         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 = Expr(:call, operate!!, add_mul, x)
+            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, is_mutable
+            return root, true
         end
         return _rewrite_generic_to_nested_op(stack, expr, add_mul)
     elseif expr.args[1] == :-

diff --git a/test/rewrite_generic.jl b/test/rewrite_generic.jl
@@ -324,6 +324,18 @@ 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(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()