Merge pull request #383 from JuliaSymbolics/myb/unflatten

Unflatten * and + in toexpr

Author: shashi

src/code.jl

@@ -7,7 +7,8 @@ export toexpr, Assignment, (←), Let, Func, DestructuredArgs, LiteralExpr,
        SpawnFetch, Multithreaded
 
 import ..SymbolicUtils
-import SymbolicUtils: @matchable, Sym, Term, istree, operation, arguments
+import SymbolicUtils: @matchable, Sym, Term, istree, operation, arguments,
+                      symtype
 
 ##== state management ==##
 
@@ -98,6 +99,22 @@ toexpr(a::Assignment, st) = :($(toexpr(a.lhs, st)) = $(toexpr(a.rhs, st)))
 
 function_to_expr(op, args, st) = nothing
 
+function function_to_expr(op::Union{typeof(*),typeof(+)}, O, st)
+    out = get(st.symbolify, O, nothing)
+    out === nothing || return out
+    args = map(Base.Fix2(toexpr, st), arguments(O))
+    if length(args) >= 3 && symtype(O) <: Number
+        x, xs = Iterators.peel(args)
+        foldl(xs, init=x) do a, b
+            Expr(:call, op, a, b)
+        end
+    else
+        expr = Expr(:call, op)
+        append!(expr.args, args)
+        expr
+    end
+end
+
 function function_to_expr(::typeof(^), O, st)
     args = arguments(O)
     if length(args) == 2 && args[2] isa Real && args[2] < 0

test/code.jl

@@ -13,15 +13,17 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
     @test toexpr(Assignment(a, b)) == :(a = b)
     @test toexpr(a ← b) == :(a = b)
     @test toexpr(a+b) == :($(+)(a, b))
+    @test toexpr(a*b*c*d*e) == :($(*)($(*)($(*)($(*)(a, b), c), d), e))
+    @test toexpr(a+b+c+d+e) == :($(+)($(+)($(+)($(+)(a, b), c), d), e))
     @test toexpr(a+b) == :($(+)(a, b))
     @test toexpr(a^b) == :($(^)(a, b))
     @test toexpr(a^2) == :($(^)(a, 2))
     @test toexpr(a^-2) == :($(^)($(inv)(a), 2))
     @test toexpr(x(t)+y(t)) == :($(+)(x(t), y(t)))
-    @test toexpr(x(t)+y(t)+x(t+1)) == :($(+)(x(t), y(t), x($(+)(1, t))))
+    @test toexpr(x(t)+y(t)+x(t+1)) == :($(+)($(+)(x(t), y(t)), x($(+)(1, t))))
     s = LazyState()
     Code.union_symbolify!(s.symbolify, [x(t), y(t)])
-    @test toexpr(x(t)+y(t)+x(t+1), s) == :($(+)(var"x(t)", var"y(t)", x($(+)(1, t))))
+    @test toexpr(x(t)+y(t)+x(t+1), s) == :($(+)($(+)(var"x(t)", var"y(t)"), x($(+)(1, t))))
 
     ex = :(let a = 3, b = $(+)(1,a)
                $(+)(a, b)
@@ -35,14 +37,14 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
 
     test_repr(toexpr(Func([x(t), x],[b ← a+2, y(t) ← b], x(t)+x(t+1)+b+y(t))),
               :(function (var"x(t)", x; b = $(+)(2, a), var"y(t)" = b)
-                    $(+)(b, var"x(t)", var"y(t)", x($(+)(1, t)))
+                    $(+)($(+)($(+)(b, var"x(t)"), var"y(t)"), x($(+)(1, t)))
                 end))
     test_repr(toexpr(Func([DestructuredArgs([x, x(t)], :state),
                            DestructuredArgs((a, b), :params)], [],
                           x(t+1) + x(t) + a  + b)),
               :(function (state, params)
                     let x = state[1], var"x(t)" = state[2], a = params[1], b = params[2]
-                        $(+)(a, b, var"x(t)", x($(+)(1, t)))
+                        $(+)($(+)($(+)(a, b), var"x(t)"), x($(+)(1, t)))
                     end
                 end))
 
@@ -81,7 +83,7 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
               :(let foo = Any[3, 3, [1, 4]],
                     var"x(t)" = foo[1], b = foo[2], c = foo[3],
                     p = c[1], q = c[2]
-                    $(+)(a, b, c, var"x(t)")
+                    $(+)($(+)($(+)(a, b), c), var"x(t)")
                 end))
 
     test_repr(toexpr(Func([DestructuredArgs([a,b],c,inds=[:a, :b])], [],
@@ -119,7 +121,7 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
 
     @test eval(toexpr(Let([a ← 1, b ← 2, arr ← @SLVector((:a, :b))(@SVector[1,2])],
                           MakeArray([a+b,a/b], arr)))) === @SLVector((:a, :b))(@SVector [3, 1/2])
-    
+
     R1 = eval(toexpr(Let([a ← 1, b ← 2, arr ← @MVector([1,2])],MakeArray([a,b,a+b,a/b], arr))))
     @test R1 == (@MVector [1, 2, 3, 1/2]) && R1 isa MVector
 
@@ -166,4 +168,3 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
         @test f(2) == 2
     end
 end
-