treat all other operations or literals as new symbols

Co-authored-by: "Yingbo Ma" <[email protected]>

Author: shashi

src/abstractalgebra.jl

@@ -1,14 +1,15 @@
+using DataStructures
+
 """
-    labels(dict, t)
+    labels!(dict, t)
 
 Find all terms that are not + and * and replace them
 with a symbol, store the symbol => term mapping in `dict`.
 """
-function labels end
+function labels! end
 
 # Turn a Term into a multivariate polynomial
-labels(dicts, t; label_terms=false) = t
-function labels(dicts, t::Sym; label_terms=false)
+function labels!(dicts, t::Sym)
     sym2term, term2sym = dicts
     if !haskey(term2sym, t)
         sym2term[t] = t
@@ -17,56 +18,71 @@ function labels(dicts, t::Sym; label_terms=false)
     return t
 end
 
-function labels(dicts, t::Term; label_terms=false)
-    tt = arguments(t)
-    if operation(t) == (*) || operation(t) == (+)
-        return Term{symtype(t)}(operation(t), map(x->labels(dicts, x;
-                                                            label_terms=label_terms), tt))
+function labels!(dicts, t)
+    if t isa Term && (operation(t) == (*) || operation(t) == (+) || operation(t) == (-))
+        tt = arguments(t)
+        return Term{symtype(t)}(operation(t), map(x->labels!(dicts, x), tt))
+    elseif t isa Integer
+        return t
     else
         sym2term, term2sym = dicts
         if haskey(term2sym, t)
             return term2sym[t]
         end
-        if label_terms
+        if t isa Term
+            tt = arguments(t)
             sym = Sym{symtype(t)}(gensym(nameof(operation(t))))
-            sym2term[sym] = Term{symtype(t)}(operation(t), map(x->labels(dicts, x;
-                                                                        label_terms=label_terms),
-                                                               tt))
-            x = term2sym[t] = sym
-
-            return x
+            sym2term[sym] = Term{symtype(t)}(operation(t),
+                                             map(x->to_mpoly(x, dicts)[1], tt))
         else
-            return Term{symtype(t)}(operation(t), map(x->labels(dicts, x; label_terms=label_terms), tt))
+            sym = Sym{symtype(t)}(gensym("literal"))
+            sym2term[sym] = t
         end
+
+        term2sym[t] = sym
+
+        return sym
     end
 end
 
 ismpoly(x) = x isa MPoly || x isa Integer
 isnonnegint(x) = x isa Integer && x >= 0
 
-function to_mpoly(t)
-    sym2term, term2sym = Dict(), Dict()
-    ls = labels((sym2term, term2sym), t)
+function to_mpoly(t, dicts=(OrderedDict(), OrderedDict()))
+    # term2sym is only used to assign the same
+    # symbol for the same term -- in other words,
+    # it does common subexpression elimination
+
+    sym2term, term2sym = dicts
+    labeled = labels!((sym2term, term2sym), t)
 
     if isempty(sym2term)
-        return t, []
+        return labeled, []
     end
 
     ks = collect(keys(sym2term))
     R, vars = PolynomialRing(ZZ, String.(nameof.(ks)))
 
-    t_poly_1 = substitute(t, term2sym, fold=false)
-    t_poly_2 = substitute(t_poly_1, Dict(ks .=> vars), fold=false)
+    t_poly = substitute(labeled, Dict(ks .=> vars), fold=false)
     rs = RuleSet([@rule(~x::ismpoly - ~y::ismpoly => ~x + -1 * (~y))
                   @acrule(~x::ismpoly + ~y::ismpoly => ~x + ~y)
                   @rule(+(~x) => ~x)
                   @acrule(~x::ismpoly * ~y::ismpoly => ~x * ~y)
                   @rule(*(~x) => ~x)
                   @rule((~x::ismpoly)^(~a::isnonnegint) => (~x)^(~a))])
-    simplify(t_poly_2, EmptyCtx(), rules=rs), Dict(Pair.(1:length(vars), ks))
+    simplify(t_poly, EmptyCtx(), rules=rs), sym2term, Dict(Pair.(1:length(vars), ks))
 end
 
-function to_term(x::MPoly, syms)
+function to_term(x::MPoly, dict, syms)
+    dict = copy(dict)
+    for (k, v) in dict
+        dict[k] = _to_term(v, dict, syms)
+    end
+    _to_term(x, dict, syms)
+end
+
+function _to_term(x::MPoly, dict, syms)
+
     function mul_coeffs(exps)
         monics = [e == 1 ? syms[i] : syms[i]^e for (i, e) in enumerate(exps) if !iszero(e)]
         if length(monics) == 1
@@ -77,20 +93,24 @@ function to_term(x::MPoly, syms)
             return Term(*, monics)
         end
     end
+
     monoms = [mul_coeffs(exponent_vector(x, i)) for i in 1:x.length]
-    if length(monoms) == 1
-        !isone(x.coeffs[1]) ?  monoms[1] * x.coeffs[1] : monoms[1]
-    elseif length(monoms) == 0
+    if length(monoms) == 0
         return 0
+    end
+    if length(monoms) == 1
+        t = !isone(x.coeffs[1]) ?  monoms[1] * x.coeffs[1] : monoms[1]
     else
-        Term(+, map((x,y)->isone(y) ? x : y*x, monoms, x.coeffs[1:length(monoms)]))
+        t = Term(+, map((x,y)->isone(y) ? x : y*x, monoms, x.coeffs[1:length(monoms)]))
     end
+
+    substitute(t, dict, fold=false)
 end
 
-to_term(x, vars) = x
+_to_term(x, dict, vars) = x
 
-function to_term(x::Term, vars)
-    Term{symtype(x)}(operation(x), to_term.(arguments(x), (vars,)))
+function _to_term(x::Term, dict, vars)
+    t=Term{symtype(x)}(operation(x), _to_term.(arguments(x), (dict,), (vars,)))
 end
 
 <ₑ(a::MPoly, b::MPoly) = false

src/simplify.jl

@@ -43,8 +43,13 @@ Applies them once if `fixpoint=false`.
 The `applyall` and `recurse` keywords are forwarded to the enclosed
 `RuleSet`, they are mainly used for internal optimization.
 """
-function simplify(x, ctx=DefaultCtx(); rules=default_rules(x, ctx), fixpoint=true, applyall=true, kwargs...)
-    if ctx isa DefaultCtx
+function simplify(x, ctx=DefaultCtx();
+                  rules=default_rules(x, ctx),
+                  fixpoint=true,
+                  applyall=true,
+                  mpoly=false,
+                  kwargs...)
+    if mpoly
         x = to_term(to_mpoly(x)...)
     end
     if fixpoint