get_degrees(): refactor

* replace if op = ... with get_degrees(op, expr) for extendability
* avoid sorted_arguments() within get_degrees()
* reduce allocations by iterating of degs in efficient way
* support /: const/expr and expr/const

Author: Alexey Stukalov

src/ordering.jl

@@ -31,28 +31,82 @@ See also `monomial_lt` and `lexlt`.
 """
 function get_degrees(expr)
     if issym(expr)
-        ((Symbol(expr),) => 1,)
+        return ((Symbol(expr),) => 1,)
     elseif iscall(expr)
-        op = operation(expr)
-        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 op == (*)
-            return mapreduce(get_degrees,
-                             (x,y)->(x...,y...,), args)
-        elseif op == (+)
-            ds = map(get_degrees, args)
-            _, idx = findmax(x->sum(last, x, init=0), ds)
-            return ds[idx]
-        elseif op == (getindex)
-            return (Tuple(map(Symbol, args)) => 1,)
-        else
-            return ((Symbol("zzzzzzz", hash(expr)),) => 1,)
+        # operation-specific degree handling
+        return get_degrees(operation(expr), expr)
+    else
+        return () # skip numbers and unsupported expressions
+    end
+end
+
+# fallback for unsupported operation
+get_degrees(::Any, expr) =
+    ((Symbol("zzzzzzz", hash(expr)),) => 1,)
+
+_getindex_symbol(arr, i) = Symbol(arr[i])
+
+function get_degrees(::typeof(getindex), expr)
+    args = arguments(expr)
+    @inbounds return (ntuple(Base.Fix1(_getindex_symbol, args), length(args)) => 1,)
+end
+
+function get_degrees(::typeof(*), expr)
+    args = arguments(expr)
+    ds = sizehint!(Vector{Any}(), length(args))
+    for arg in args
+        degs = get_degrees(arg)
+        append!(ds, degs)
+    end
+    return sort!(ds)
+end
+
+function get_degrees(::typeof(+), expr)
+    # among the terms find the best in terms of monomial_lt
+    sel_degs = ()
+    sel_degsum = 0
+    for arg in arguments(expr)
+        degs = get_degrees(arg)
+        degsum = sum(last, degs, init=0)
+        if (sel_degs == ()) || (degsum > sel_degsum) ||
+           (degsum == sel_degsum && lexlt(degs, sel_degs))
+            sel_degs, sel_degsum = degs, degsum
+        end
+    end
+    return sel_degs
+end
+
+function get_degrees(::typeof(^), expr)
+    base_expr, pow_expr = arguments(expr)
+    if pow_expr isa Number
+        @inbounds degs = map(get_degrees(base_expr)) do (base, pow)
+            (base => pow * pow_expr)
+        end
+        if pow_expr < 0 && length(degs) > 1
+            # fix the order after the powers were negated
+            isa(degs, AbstractVector) || (degs = collect(degs))
+            sort!(degs)
+        end
+        return degs
+    else
+        # expression in the power argument is not supported
+        return get_degrees(nothing, expr)
+    end
+end
+
+function get_degrees(::typeof(/), expr)
+    nom_expr, denom_expr = arguments(expr)
+    if denom_expr isa Number # constant denominator
+        return get_degrees(nom_expr)
+    elseif nom_expr isa Number # constant nominator
+        @inbounds degs = map(get_degrees(denom_expr)) do (base, pow)
+            (base => -pow)
         end
+        isa(degs, AbstractVector) || (degs = collect(degs))
+        return sort!(degs)
     else
-        return ()
+        # TODO expressions in both nom and denom are not yet supported
+        return get_degrees(nothing, expr)
     end
 end
 

test/basics.jl

@@ -221,6 +221,8 @@ end
     b1, b3, d1, d2 = get(b,1),get(b,3), get(d,1), get(d,2)
     @test repr(a + b3 + b1 + d2 + c) == "a + b[1] + b[3] + c + d[2]"
     @test repr(expand((c + b3 - d1)^3)) == "b[3]^3 + 3(b[3]^2)*c - 3(b[3]^2)*d[1] + 3b[3]*(c^2) - 6b[3]*c*d[1] + 3b[3]*(d[1]^2) + c^3 - 3(c^2)*d[1] + 3c*(d[1]^2) - (d[1]^3)"
+    # test negative powers sorting
+    @test repr((b3^2)^(-2) + a^(-3) + (c*d1)^(-2)) == "1 / (b[3]^4) + 1 / ((c^2)*(d[1]^2)) + 1 / (a^3)"
 end
 
 @testset "inspect" begin