feat: add fast in-place Polynomial + constant implementation

Author: AayushSabharwal

src/operators.jl

@@ -92,6 +92,47 @@ function MA.operate_to!(
     return output
 end
 
+"""
+    _constant_term_idx(p::Polynomial)
+
+Return the index of the constant term in `p` according to its monomial ordering.
+"""
+_constant_term_idx(p::Polynomial{V, M, T}) where {V, M <: Reverse, T} = lastindex(p.x)
+_constant_term_idx(p::Polynomial) = firstindex(p.x)
+
+"""
+    _insert_constant_term!(p::Polynomial)
+
+Insert a constant (degree 0) term into polynomial `p` at the appropriate position for the
+monomial ordering of `p`. Does not check if a constant term already exists.
+"""
+function _insert_constant_term!(p::Polynomial{V, M, T}) where {V, M <: Reverse, T}
+    push!(MP.coefficients(p), zero(T))
+    push!(MP.monomials(p).Z, zeros(Int, length(MP.variables(p))))
+    return p
+end
+
+function _insert_constant_term!(p::Polynomial{V, M, T}) where {V, M, T}
+    insert!(MP.coefficients(p), 1, zero(T))
+    insert!(MP.monomials(p).Z, 1, zeros(Int, length(MP.variables(p))))
+    return p
+end
+
+function MA.operate!(op::Union{typeof(+), typeof(-)}, p::Polynomial{V, M, T}, x::T) where {V, M, T}
+    c_idx = _constant_term_idx(p)
+    if MP.nterms(p) == 0 || !MP.isconstant(MP.terms(p)[c_idx])
+        _insert_constant_term!(p)
+        c_idx = _constant_term_idx(p)
+    end
+    coeffs = MP.coefficients(p)
+    coeffs[c_idx] = op(coeffs[c_idx], x)
+    if iszero(coeffs[c_idx])
+        deleteat!(coeffs, c_idx)
+        deleteat!(MP.monomials(p), c_idx)
+    end
+    return p
+end
+
 function MA.operate!(
     op::Union{typeof(+),typeof(-)},
     p::Polynomial,

test/mutable_arithmetics.jl

@@ -35,3 +35,42 @@ using DynamicPolynomials
         @test q == x + y + 1
     end
 end
+
+@testset "Fast path cases: $ord" for ord in [
+        InverseLexOrder,
+        LexOrder,
+        Graded{InverseLexOrder},
+        Graded{LexOrder},
+        Reverse{InverseLexOrder},
+        Reverse{LexOrder},
+        Graded{Reverse{InverseLexOrder}},
+        Graded{Reverse{LexOrder}},
+        Reverse{Graded{InverseLexOrder}},
+        Reverse{Graded{LexOrder}},
+    ]
+    @polyvar x y z monomial_order=ord
+
+    # allocation tests vary between Julia versions, so they're upper bounds
+    @testset "Polynomial + constant" begin
+        poly = 2 * x^2 + 3 * x * y + z * y^2
+        poly2 = copy(poly)
+        result = poly + 2
+        MA.operate!(+, poly, 2)
+        @test isequal(poly, result)
+        # down from 576 using the generic method
+        @test (@allocated MA.operate!(+, poly2, 2)) <= 272
+        # subsequent additions don't allocate
+        @test (@allocated MA.operate!(+, poly2, 2)) == 0
+
+        # also test `-`
+        poly = 2 * x^2 + 3 * x * y + z * y^2
+        poly2 = copy(poly)
+        result = poly - 2
+        MA.operate!(-, poly, 2)
+        @test isequal(poly, result)
+        # down from 576 using the generic method
+        @test (@allocated MA.operate!(-, poly2, 2)) <= 272
+        # subsequent additions don't allocate
+        @test (@allocated MA.operate!(-, poly2, 2)) == 0
+    end
+end