more improvements

Author: oameye

benchmark/baseline.json

@@ -1 +1,5 @@
-pages = ["Home" => "index.md", "API" => "API.md"]
+pages = [
+    "Home" => "index.md",
+    "API" => "API.md",
+    "Developer" => ["Symbolic Rewriting" => "dev/symbolic_rewriting.md"],
+]

docs/pages.jl

@@ -0,0 +1,102 @@
+# Symbolic Rewriting: Design Decisions
+
+This document records what was tried when evaluating SymbolicUtils.jl's `@rule`-based
+rewriting system as a replacement for QuestBase's manual tree traversal.
+
+## What We Tried
+
+We attempted to replace the manual `isadd`/`ismul`/`isdiv` dispatch in four functions
+with `@rule` / `@acrule` + `Postwalk`:
+
+### `expand_exp_power`
+
+```julia
+# Before: manual recursion (14 μs)
+function expand_exp_power(expr::BasicSymbolic)
+    if isadd(expr)
+        sum(expand_exp_power(arg) for arg in arguments(expr))
+    elseif ismul(expr)
+        prod(expand_exp_power(arg) for arg in arguments(expr))
+    elseif ispow(expr) && isexp(arguments(expr)[1])
+        exp(arguments(arguments(expr)[1])[1] * arguments(expr)[2])
+    else
+        expr
+    end
+end
+
+# After: single rule + Postwalk (28 μs, 1.9x slower)
+const _expand_exp_power_rw = Postwalk(PassThrough(@rule exp(~x)^~n => exp(~x * ~n)))
+expand_exp_power(expr::BasicSymbolic) = _expand_exp_power_rw(expr)
+```
+
+### `simplify_exp_products`
+
+```julia
+# Before: manual findall + index manipulation (10 μs)
+# (see src/Symbolics/exponentials.jl)
+
+# After: @acrule + Fixpoint + Postwalk (71 μs, 6.8x slower)
+const _simplify_exp_products_rw = Fixpoint(
+    Postwalk(
+        PassThrough(
+            Chain((
+                @acrule(exp(~a) * exp(~b) => exp(~a + ~b)),
+                @rule(exp(~x)^~n => exp(~x * ~n)),
+                @rule(exp(~x::is_literal_number) => begin
+                    iszero(unwrap_const(~x)) ? 1 : exp(~x)
+                end),
+            )),
+        ),
+    ),
+)
+```
+
+The `Fixpoint` is needed because combining two exps can create a new `exp(x)^2` (via
+hash-consing auto-simplification), which then needs the power rule, which may enable
+further combining.
+
+### `expand_fraction`
+
+```julia
+# Before: manual _apply_termwise dispatch (25 μs)
+
+# After: single rule + Postwalk (37 μs, 1.5x slower)
+const _expand_fraction_rw = Postwalk(
+    PassThrough(@rule (+(~~xs)) / ~c => sum(x / ~c for x in ~~xs))
+)
+```
+
+### `exp_to_trig`
+
+```julia
+# Before: manual _apply_termwise + inline conversion (0.01 μs on no-op input)
+
+# After: Postwalk with custom node function (27 μs on no-op input, 2500x slower)
+exp_to_trig(x::BasicSymbolic) = Postwalk(PassThrough(_exp_to_trig_node))(x)
+```
+
+The extreme regression on no-op input (expression with no `exp` nodes) happens because
+`Postwalk` visits every node in the tree. The manual code sees `isadd` at the top level,
+recurses into children, sees they're all `cos`/`sin` terms (not `exp`), and returns
+immediately.
+
+## Why It's Slower
+
+For QuestBase's typical expressions (5-50 nodes):
+
+1. **`Postwalk` visits every node** — even when only a few nodes match. Manual code
+   skips irrelevant subtrees via `_apply_termwise`.
+2. **Pattern matching overhead** — `@rule` builds closures with continuation-passing
+   style and `ImmutableDict` bindings. For a simple `ispow && isexp` check, this is
+   much heavier than two boolean tests.
+3. **`Fixpoint` re-traverses** — When rules interact, `Fixpoint(Postwalk(...))` walks
+   the full tree again. Manual code handles interactions in a single pass.
+4. **`@acrule` tries permutations** — For `exp(a)*exp(b)`, it tries all argument
+   orderings up to `COMM_CHECKS_LIMIT` (default 10).
+
+The rule system is designed for Symbolics.jl's simplifier, which applies hundreds of
+rules to large expression trees. The overhead is amortized there but dominates for
+QuestBase's small, targeted transformations.
+
+`Postwalk` is still used where it works well: `expand_all` and `add_div`, where the
+rewriter needs to visit all nodes anyway.

docs/src/dev/symbolic_rewriting.md

@@ -14,7 +14,7 @@ function expand_fraction(x::BasicSymbolic)
         _apply_termwise(expand_fraction, x)
     elseif isdiv(x)
         args = arguments(x)
-        sum([arg / args[2] for arg in arguments(args[1])])
+        sum(arg / args[2] for arg in arguments(args[1]))
     else
         x
     end

src/Symbolics/Symbolics_utils.jl

@@ -1,52 +1,54 @@
-expand_exp_power(expr::Num) = expand_exp_power(unwrap(expr))
-simplify_exp_products(x::Num) = simplify_exp_products(unwrap(x))
-
 "Returns true if expr is an exponential"
 isexp(expr) = isterm(expr) && operation(expr) === exp
 
-"Expand powers of exponential such that exp(x)^n => exp(x*n) "
+"Expand powers of exponential such that exp(x)^n => exp(x*n)"
 function expand_exp_power(expr::BasicSymbolic)
     if isadd(expr)
         sum(expand_exp_power(arg) for arg in arguments(expr))
     elseif ismul(expr)
         prod(expand_exp_power(arg) for arg in arguments(expr))
+    elseif ispow(expr) && isexp(arguments(expr)[1])
+        exp(arguments(arguments(expr)[1])[1] * arguments(expr)[2])
     else
-        if ispow(expr) && isexp(arguments(expr)[1])
-            exp(arguments(arguments(expr)[1])[1] * arguments(expr)[2])
-        else
-            expr
-        end
+        expr
     end
 end
+expand_exp_power(expr::Num) = expand_exp_power(unwrap(expr))
 expand_exp_power(expr) = expr
 
-"Simplify products of exponentials such that exp(a)*exp(b) => exp(a+b)
-This is included in SymbolicUtils as of 17.0 but the method here avoid other simplify calls"
+"Simplify products of exponentials such that exp(a)*exp(b) => exp(a+b).
+Also expands exp(x)^n => exp(x*n) and simplifies exp(0) => 1."
 function simplify_exp_products(expr::BasicSymbolic)
     if isadd(expr)
         _apply_termwise(simplify_exp_products, expr)
     elseif isdiv(expr)
         _apply_termwise(simplify_exp_products, expr)
     elseif ismul(expr)
-        simplify_exp_products_mul(expr)
+        _simplify_exp_products_mul(expr)
+    elseif ispow(expr) && isexp(arguments(expr)[1])
+        # exp(x)^n => exp(x*n) — fixes bug where exp powers were left unexpanded
+        exp(arguments(arguments(expr)[1])[1] * arguments(expr)[2])
     else
         expr
     end
 end
+
+function _simplify_exp_products_mul(expr)
+    args = arguments(expr)
+    ind = findall(isexp, args)
+    rest_ind = setdiff(1:length(args), ind)
+    rest = isempty(rest_ind) ? 1 : prod(args[rest_ind])
+    total = isempty(ind) ? 0 : sum(arguments(args[i])[1] for i in ind)
+    if is_literal_number(total)
+        iszero(unwrap_const(total)) && return rest
+    end
+    return isempty(ind) ? rest : rest * exp(total)
+end
+
+simplify_exp_products(x::Num) = simplify_exp_products(unwrap(x))
 function simplify_exp_products(x::Complex{Num})
     return Complex{Num}(
         simplify_exp_products(unwrap(x.re)), simplify_exp_products(unwrap(x.im))
     )
 end
-function simplify_exp_products_mul(expr)
-    ind = findall(x -> isexp(x), arguments(expr))
-    rest_ind = setdiff(1:length(arguments(expr)), ind)
-    rest = isempty(rest_ind) ? 1 : prod(arguments(expr)[rest_ind])
-    total = isempty(ind) ? 0 : sum(getindex.(arguments.(arguments(expr)[ind]), 1))
-    if is_literal_number(total)
-        (iszero(unwrap_const(total)) && return rest)
-    else
-        return rest * exp(total)
-    end
-end
 simplify_exp_products(x) = x

src/Symbolics/exponentials.jl

@@ -191,6 +191,54 @@ function trig_to_exp(x::BasicSymbolic)
     return Symbolics.substitute(x, Dict(rules))
 end
 
+"Convert a single exp(arg) node to trig form with sign normalization.
+Returns `nothing` if the node is not an exp call (for use with PassThrough)."
+function _exp_to_trig_node(x::BasicSymbolic)
+    !(isterm(x) && operation(x) === exp) && return nothing
+
+    arg = first(arguments(x))
+
+    # exp(0) = 1: handle literal zero argument
+    arg_val = arg isa BasicSymbolic ? unwrap_const(arg) : arg
+    if arg_val isa Number && iszero(arg_val)
+        return 1
+    end
+
+    trigarg = Symbolics.expand(-im * arg) # the argument of the to-be trig function
+    trigarg = simplify_complex(trigarg)
+
+    # After expansion and simplification, check if trigarg is zero (exp(0) = 1)
+    ta_val = trigarg isa BasicSymbolic ? unwrap_const(trigarg) : trigarg
+    if ta_val isa Number && iszero(ta_val)
+        return 1
+    end
+
+    # put arguments of trigs into a standard form such that sin(x) = -sin(-x), cos(x) = cos(-x) are recognized
+    if isadd(trigarg)
+        first_symbol = minimum(
+            cat(
+                string.(sorted_arguments(trigarg)),
+                string.(sorted_arguments(-trigarg));
+                dims=1,
+            ),
+        )
+
+        # put trigarg => -trigarg the lowest alphabetic argument of trigarg is lower than that of -trigarg
+        # this is a meaningless key but gives unique signs to all sums
+        is_first = minimum(string.(sorted_arguments(trigarg))) == first_symbol
+        return if is_first
+            cos(-trigarg) - im * sin(-trigarg)
+        else
+            cos(trigarg) + im * sin(trigarg)
+        end
+    end
+    return if ismul(trigarg) && _has_negative_coefficient(trigarg)
+        cos(-trigarg) - im * sin(-trigarg)
+    else
+        cos(trigarg) + im * sin(trigarg)
+    end
+end
+
 """
     exp_to_trig(x::BasicSymbolic)
     exp_to_trig(x)
@@ -210,51 +258,9 @@ trigonometric arguments for consistent simplification.
 function exp_to_trig(x::BasicSymbolic)
     if isadd(x) || isdiv(x) || ismul(x)
         return _apply_termwise(exp_to_trig, x)
-    elseif isterm(x) && operation(x) === exp
-        arg = first(arguments(x))
-
-        # exp(0) = 1: handle literal zero argument
-        arg_val = arg isa BasicSymbolic ? unwrap_const(arg) : arg
-        if arg_val isa Number && iszero(arg_val)
-            return 1
-        end
-
-        trigarg = Symbolics.expand(-im * arg) # the argument of the to-be trig function
-        trigarg = simplify_complex(trigarg)
-
-        # After expansion and simplification, check if trigarg is zero (exp(0) = 1)
-        ta_val = trigarg isa BasicSymbolic ? unwrap_const(trigarg) : trigarg
-        if ta_val isa Number && iszero(ta_val)
-            return 1
-        end
-
-        # put arguments of trigs into a standard form such that sin(x) = -sin(-x), cos(x) = cos(-x) are recognized
-        if isadd(trigarg)
-            first_symbol = minimum(
-                cat(
-                    string.(sorted_arguments(trigarg)),
-                    string.(sorted_arguments(-trigarg));
-                    dims=1,
-                ),
-            )
-
-            # put trigarg => -trigarg the lowest alphabetic argument of trigarg is lower than that of -trigarg
-            # this is a meaningless key but gives unique signs to all sums
-            is_first = minimum(string.(sorted_arguments(trigarg))) == first_symbol
-            return if is_first
-                cos(-trigarg) - im * sin(-trigarg)
-            else
-                cos(trigarg) + im * sin(trigarg)
-            end
-        end
-        return if ismul(trigarg) && _has_negative_coefficient(trigarg)
-            cos(-trigarg) - im * sin(-trigarg)
-        else
-            cos(trigarg) + im * sin(trigarg)
-        end
-    else
-        return x
     end
+    result = _exp_to_trig_node(x)
+    return isnothing(result) ? x : result
 end
 
 "Check if a Mul expression has a negative leading coefficient"