Merge pull request #615 from JuliaSymbolics/b/613-default-arguments-to-unsorted_arguments-to-accelerate-term-traversal

Optimize `arguments` function by removing sorting

Author: ChrisRackauckas

src/SymbolicUtils.jl

@@ -21,7 +21,7 @@ import TermInterface: iscall, isexpr, issym, symtype, head, children,
 
 const istree = iscall
 Base.@deprecate_binding istree iscall
-export istree, operation, arguments, unsorted_arguments, similarterm, iscall
+export istree, operation, arguments, sorted_arguments, similarterm, iscall
 # Sym, Term,
 # Add, Mul and Pow
 include("types.jl")

src/code.jl

@@ -9,7 +9,7 @@ export toexpr, Assignment, (←), Let, Func, DestructuredArgs, LiteralExpr,
 import ..SymbolicUtils
 import ..SymbolicUtils.Rewriters
 import SymbolicUtils: @matchable, BasicSymbolic, Sym, Term, iscall, operation, arguments, issym,
-                      symtype, similarterm, unsorted_arguments, metadata, isterm, term, maketerm
+                      symtype, similarterm, sorted_arguments, metadata, isterm, term, maketerm
 
 ##== state management ==##
 
@@ -124,7 +124,7 @@ end
 function function_to_expr(op::Union{typeof(*),typeof(+)}, O, st)
     out = get(st.rewrites, O, nothing)
     out === nothing || return out
-    args = map(Base.Fix2(toexpr, st), arguments(O))
+    args = map(Base.Fix2(toexpr, st), sorted_arguments(O))
     if length(args) >= 3 && symtype(O) <: Number
         x, xs = Iterators.peel(args)
         foldl(xs, init=x) do a, b
@@ -744,7 +744,7 @@ end
 function cse_state!(state, t)
     !iscall(t) && return t
     state[t] = Base.get(state, t, 0) + 1
-    foreach(x->cse_state!(state, x), unsorted_arguments(t))
+    foreach(x->cse_state!(state, x), arguments(t))
 end
 
 function cse_block!(assignments, counter, names, name, state, x)
@@ -759,7 +759,7 @@ function cse_block!(assignments, counter, names, name, state, x)
             return sym
         end
     elseif iscall(x)
-        args = map(a->cse_block!(assignments, counter, names, name, state,a), unsorted_arguments(x))
+        args = map(a->cse_block!(assignments, counter, names, name, state,a), arguments(x))
         if isterm(x)
             return term(operation(x), args...)
         else

src/inspect.jl

@@ -26,8 +26,16 @@ function AbstractTrees.nodevalue(x::BasicSymbolic)
     Text(str)
 end
 
+"""
+$(TYPEDSIGNATURES)
+
+Return the children of the symbolic expression `x`, sorted by their order in
+the expression.
+
+This function is used internally for printing via AbstractTrees.
+"""
 function AbstractTrees.children(x::Symbolic)
-    iscall(x) ? arguments(x) : isexpr(x) ? children(x) : ()
+    iscall(x) ? sorted_arguments(x) : isexpr(x) ? sorted_children(x) : ()
 end
 
 """

src/interface.jl

@@ -36,22 +36,22 @@ is the function being called.
 function operation end
 
 """
-  arguments(x)
+    sorted_arguments(x)
 
 Get the arguments of `x`, must be defined if `iscall(x)` is `true`.
 """
-function arguments end
+function sorted_arguments end
 
 """
-  unsorted_arguments(x::T)
+    sorted_arguments(x::T)
 
 If x is a term satisfying `iscall(x)` and your term type `T` provides
 an optimized implementation for storing the arguments, this function can
 be used to retrieve the arguments when the order of arguments does not matter 
 but the speed of the operation does.
 """
-unsorted_arguments(x) = arguments(x)
-arity(x) = length(unsorted_arguments(x))
+function arguments end
+arity(x) = length(arguments(x))
 
 """
   metadata(x)

src/ordering.jl

@@ -14,28 +14,31 @@
 <ₑ(a::T, b::S) where{T,S} = T<S
 <ₑ(a::T, b::T) where{T} = a < b
 
+"""
+$(SIGNATURES)
 
-###### A variation on degree lexicographic order ########
-# find symbols and their corresponding degrees
+Internal function used for printing symbolic expressions. This function determines
+the degrees of symbols within a given expression, implementing a variation on 
+degree lexicographic order.
+"""
 function get_degrees(expr)
     if issym(expr)
         ((Symbol(expr),) => 1,)
     elseif iscall(expr)
         op = operation(expr)
-        args = arguments(expr)
-        if operation(expr) == (^) && args[2] isa Number
+        args = sorted_arguments(expr)
+        if op == (^) && args[2] isa Number
             return map(get_degrees(args[1])) do (base, pow)
                 (base => pow * args[2])
             end
-        elseif operation(expr) == (*)
+        elseif op == (*)
             return mapreduce(get_degrees,
                              (x,y)->(x...,y...,), args)
-        elseif operation(expr) == (+)
+        elseif op == (+)
             ds = map(get_degrees, args)
             _, idx = findmax(x->sum(last.(x), init=0), ds)
             return ds[idx]
-        elseif operation(expr) == (getindex)
-            args = arguments(expr)
+        elseif op == (getindex)
             return ((Symbol.(args)...,) => 1,)
         else
             return ((Symbol("zzzzzzz", hash(expr)),) => 1,)
@@ -62,7 +65,7 @@ function lexlt(degs1, degs2)
     return false # they are equal
 end
 
-_arglen(a) = iscall(a) ? length(unsorted_arguments(a)) : 0
+_arglen(a) = iscall(a) ? length(arguments(a)) : 0
 
 function <ₑ(a::Tuple, b::Tuple)
     for (x, y) in zip(a, b)

src/polyform.jl

@@ -231,6 +231,9 @@ function arguments(x::PolyForm{T}) where {T}
                  PolyForm{T}(t, x.pvar2sym, x.sym2term, nothing)) for t in ts]
     end
 end
+
+sorted_arguments(x::PolyForm) = arguments(x)
+
 children(x::PolyForm) = [operation(x); arguments(x)]
 
 Base.show(io::IO, x::PolyForm) = show_term(io, x)
@@ -344,7 +347,7 @@ end
 
 function add_with_div(x, flatten=true)
     (!iscall(x) || operation(x) != (+)) && return x
-    aa = unsorted_arguments(x)
+    aa = arguments(x)
     !any(a->isdiv(a), aa) && return x # no rewrite necessary
 
     divs = filter(a->isdiv(a), aa)
@@ -382,12 +385,12 @@ end
 
 function needs_div_rules(x)
     (isdiv(x) && !(x.num isa Number) && !(x.den isa Number)) ||
-    (iscall(x) && operation(x) === (+) && count(has_div, unsorted_arguments(x)) > 1) ||
-    (iscall(x) && any(needs_div_rules, unsorted_arguments(x)))
+    (iscall(x) && operation(x) === (+) && count(has_div, arguments(x)) > 1) ||
+    (iscall(x) && any(needs_div_rules, arguments(x)))
 end
 
 function has_div(x)
-    return isdiv(x) || (iscall(x) && any(has_div, unsorted_arguments(x)))
+    return isdiv(x) || (iscall(x) && any(has_div, arguments(x)))
 end
 
 flatten_pows(xs) = map(xs) do x
@@ -415,8 +418,8 @@ Has optimized processes for `Mul` and `Pow` terms.
 function quick_cancel(d)
     if ispow(d) && isdiv(d.base)
         return quick_cancel((d.base.num^d.exp) / (d.base.den^d.exp))
-    elseif ismul(d) && any(isdiv, unsorted_arguments(d))
-        return prod(unsorted_arguments(d))
+    elseif ismul(d) && any(isdiv, arguments(d))
+        return prod(arguments(d))
     elseif isdiv(d)
         num, den = quick_cancel(d.num, d.den)
         return Div(num, den)

src/rewriters.jl

@@ -33,7 +33,7 @@ module Rewriters
 using SymbolicUtils: @timer
 using TermInterface
 
-import SymbolicUtils: iscall, operation, arguments, unsorted_arguments, metadata, node_count, _promote_symtype
+import SymbolicUtils: iscall, operation, arguments, sorted_arguments, metadata, node_count, _promote_symtype
 export Empty, IfElse, If, Chain, RestartedChain, Fixpoint, Postwalk, Prewalk, PassThrough
 
 # Cache of printed rules to speed up @timer
@@ -221,7 +221,7 @@ function (p::Walk{ord, C, F, false})(x) where {ord, C, F}
 
         if iscall(x)
             x = p.maketerm(x, operation(x), map(PassThrough(p),
-                            unsorted_arguments(x)), metadata=metadata(x))
+                            arguments(x)), metadata=metadata(x))
         end
 
         return ord === :post ? p.rw(x) : x

src/rule.jl

@@ -399,7 +399,7 @@ function (acr::ACRule)(term)
         end
 
         T = symtype(term)
-        args = unsorted_arguments(term)
+        args = arguments(term)
 
         itr = acr.sets(eachindex(args), acr.arity)
 

src/simplify.jl

@@ -45,6 +45,6 @@ end
 
 has_operation(x, op) = (iscall(x) && (operation(x) == op ||
                                       any(a->has_operation(a, op),
-                                          unsorted_arguments(x))))
+                                          arguments(x))))
 
 Base.@deprecate simplify(x, ctx; kwargs...)  simplify(x; rewriter=ctx, kwargs...)

src/substitute.jl

@@ -20,15 +20,15 @@ function substitute(expr, dict; fold=true)
         op = substitute(operation(expr), dict; fold=fold)
         if fold
             canfold = !(op isa Symbolic)
-            args = map(unsorted_arguments(expr)) do x
+            args = map(arguments(expr)) do x
                 x′ = substitute(x, dict; fold=fold)
                 canfold = canfold && !(x′ isa Symbolic)
                 x′
             end
             canfold && return op(args...)
             args
         else
-            args = map(x->substitute(x, dict, fold=fold), unsorted_arguments(expr))
+            args = map(x->substitute(x, dict, fold=fold), arguments(expr))
         end
 
         maketerm(typeof(expr),