feat: add fast in-place Polynomial + Variable implementation

Author: AayushSabharwal

src/operators.jl

@@ -129,6 +129,33 @@ function MA.operate!(op::Union{typeof(+), typeof(-)}, p::Polynomial{V, M, T}, x:
     return p
 end
 
+function MA.operate!(op::Union{typeof(+), typeof(-)}, p::Polynomial{V, M, T}, x::Variable{V, M}) where {V, M, T}
+    vars = MP.variables(p)
+    idx = searchsortedfirst(vars, x; rev = true)
+    monos = MP.monomials(p)
+    if idx > length(vars) || !isequal(vars[idx], x)
+        for mono in monos
+            insert!(MP.exponents(mono), idx, 0)
+        end
+        insert!(vars, idx, x)
+    end
+    mono = Monomial{V, M}(vars, zeros(Int, length(vars)))
+    mono.z[idx] = 1
+    idx = searchsortedfirst(monos, mono)
+    coeffs = MP.coefficients(p)
+    N = MP.nterms(p)
+    if idx > N || !isequal(monos[idx], mono)
+        insert!(monos.Z, idx, MP.exponents(mono))
+        insert!(coeffs, idx, zero(T))
+    end
+    coeffs[idx] = op(coeffs[idx], one(T))
+    if iszero(coeffs[idx])
+        deleteat!(coeffs, idx)
+        deleteat!(monos, idx)
+    end
+    return p
+end
+
 function MA.operate!(
     op::Union{typeof(+),typeof(-)},
     p::Polynomial,

test/mutable_arithmetics.jl

@@ -73,4 +73,29 @@ end
         # subsequent additions don't allocate
         @test (@allocated MA.operate!(-, poly2, 2)) == 0
     end
+
+    @testset "Polynomial + Variable" begin
+        poly = 2 * x^2 + 3 * x * y + z * y^2
+        poly2 = copy(poly)
+        result = poly + x
+        MA.operate!(+, poly, x)
+        @test isequal(poly, result)
+        # down from 18752 using the generic method
+        # 368 or 304 depending on ordering
+        @test (@allocated MA.operate!(+, poly2, x)) <= 368
+        # down from 1904 using the generic method
+        @test (@allocated MA.operate!(+, poly2, x)) <= 112
+
+        # also test `-`
+        poly = 2 * x^2 + 3 * x * y + z * y^2
+        poly2 = copy(poly)
+        result = poly - x
+        MA.operate!(-, poly, x)
+        @test isequal(poly, result)
+        # down from 18752 using the generic method
+        # 368 or 304 depending on ordering
+        @test (@allocated MA.operate!(-, poly2, x)) <= 368
+        # down from 1904 using the generic method
+        @test (@allocated MA.operate!(-, poly2, x)) <= 112
+    end
 end