Add sort kwarg in arguments to avoid extra computation (#337)

Author: YingboMa

src/api.jl

@@ -57,21 +57,22 @@ function substitute(expr, dict; fold=true)
     haskey(dict, expr) && return dict[expr]
 
     if istree(expr)
+        op = substitute(operation(expr), dict; fold=fold)
         if fold
-            canfold = !(operation(expr) isa Symbolic)
-            args = map(arguments(expr)) do x
+            canfold = !(op isa Symbolic)
+            args = map(unsorted_arguments(expr)) do x
                 x′ = substitute(x, dict; fold=fold)
                 canfold = canfold && !(x′ isa Symbolic)
                 x′
             end
-            canfold && return operation(expr)(args...)
+            canfold && return op(args...)
             args
         else
-            args = map(x->substitute(x, dict, fold=fold), arguments(expr))
+            args = map(x->substitute(x, dict, fold=fold), unsorted_arguments(expr))
         end
 
         similarterm(expr,
-                    substitute(operation(expr), dict, fold=fold),
+                    op,
                     args,
                     symtype(expr),
                     metadata=metadata(expr))

src/code.jl

@@ -168,7 +168,7 @@ function get_assignments(d::DestructuredArgs, st)
         ex = (i isa Symbol ? :($name.$i) : :($name[$i]))
         ex = d.inbounds ? :(@inbounds($ex)) : ex
         a ← ex
-    end   
+    end
 end
 
 @matchable struct Let

src/types.jl

@@ -353,6 +353,7 @@ end
 
 operation(x::Term) = getfield(x, :f)
 
+unsorted_arguments(x) = arguments(x)
 arguments(x::Term) = getfield(x, :arguments)
 
 function Base.isequal(t1::Term, t2::Term)
@@ -632,6 +633,11 @@ istree(a::Add) = true
 
 operation(a::Add) = +
 
+function unsorted_arguments(a::Add)
+    args = [v*k for (k,v) in a.dict]
+    iszero(a.coeff) ? args : vcat(a.coeff, args)
+end
+
 function arguments(a::Add)
     a.sorted_args_cache[] !== nothing && return a.sorted_args_cache[]
     args = sort!([v*k for (k,v) in a.dict], lt=<ₑ)
@@ -777,6 +783,11 @@ operation(a::Mul) = *
 
 unstable_pow(a, b) = a isa Integer && b isa Integer ? (a//1) ^ b : a ^ b
 
+function unsorted_arguments(a::Mul)
+    args = [unstable_pow(k, v) for (k,v) in a.dict]
+    isone(a.coeff) ? args : vcat(a.coeff, args)
+end
+
 function arguments(a::Mul)
     a.sorted_args_cache[] !== nothing && return a.sorted_args_cache[]
     args = sort!([unstable_pow(k, v) for (k,v) in a.dict], lt=<ₑ)

test/code.jl

@@ -150,13 +150,13 @@ test_repr(a, b) = @test repr(Base.remove_linenums!(a)) == repr(Base.remove_linen
               :((a,b,$(+)(a,b))))
 
     @test SpawnFetch{Multithreaded}([()->1,()->2],vcat)|>toexpr|>eval == [1,2]
-    @test @elapsed(SpawnFetch{Multithreaded}([:(()->sleep(.6)),
+    @test @elapsed(SpawnFetch{Multithreaded}([:(()->sleep(2)),
                                               Func([:x],
                                                    [],
                                                    :(sleep(x)))],
                                              [(),
-                                              (0.6,)],
-                                             vcat)|>toexpr|>eval) < 1.1
+                                              (2,)],
+                                             vcat)|>toexpr|>eval) < 3
 
     let
         @syms a b