Zero-allocation terms/coefficients/monomials for AbstractTermLike (#210)

* Zero-allocation terms/coefficients/monomials for AbstractTermLike

* Fix

* Fix

Author: blegat

src/MultivariatePolynomials.jl

@@ -63,6 +63,7 @@ abstract type AbstractPolynomial{T} <: AbstractPolynomialLike{T} end
 const APL{T} = AbstractPolynomialLike{T}
 
 include("zip.jl")
+include("lazy_iterators.jl")
 
 include("variable.jl")
 include("monomial.jl")
@@ -91,7 +92,6 @@ include("chain_rules.jl")
 
 include("default_term.jl")
 include("sequences.jl")
-include("lazy_iterators.jl")
 include("default_polynomial.jl")
 
 end # module

src/default_polynomial.jl

@@ -150,13 +150,13 @@ end
 Base.copy(p::VectorPolynomial) = MA.mutable_copy(p)
 
 function grlex end
-function MA.operate!(op::Union{typeof(+), typeof(-)}, p::Polynomial{T,TT}, q::Polynomial) where {T,TT}
+function MA.operate!(op::Union{typeof(+), typeof(-)}, p::Polynomial{T,TT}, q::Union{Polynomial,AbstractTermLike}) where {T,TT}
     get1 = let p=p
         i -> p.terms[i]
     end
     get2 = let p=p
         i -> begin
-            t = q.terms[i]
+            t = terms(q)[i]
             TT(MA.scaling_convert(T, MA.operate(op, coefficient(t))), monomial(t))
         end
     end
@@ -173,15 +173,15 @@ function MA.operate!(op::Union{typeof(+), typeof(-)}, p::Polynomial{T,TT}, q::Po
             if t isa Int && j isa Int
                 t = get1(t)
             end
-            grlex(monomial(q.terms[j]), monomial(t))
+            grlex(monomials(q)[j], monomial(t))
         end
     end
     combine = let p=p, q=q
         (i, j) -> begin
             if i isa Int
-                p.terms[i] = Term(MA.operate!!(op, coefficient(p.terms[i]), coefficient(q.terms[j])), monomial(p.terms[i]))
+                p.terms[i] = Term(MA.operate!!(op, coefficient(p.terms[i]), coefficients(q)[j]), monomial(p.terms[i]))
             else
-                typeof(i)(MA.operate!!(op, coefficient(i), coefficient(q.terms[j])), monomial(i))
+                typeof(i)(MA.operate!!(op, coefficient(i), coefficients(q)[j]), monomial(i))
             end
         end
     end

src/lazy_iterators.jl

@@ -1,3 +1,26 @@
+# Allows `terms(::AbstractTerm)` to not allocate and be type stable
+struct OneOrZeroElementVector{T} <: AbstractVector{T}
+    empty::Bool
+    el::T
+end
+
+Base.size(it::OneOrZeroElementVector) = (length(it),)
+
+Base.length(it::OneOrZeroElementVector) = Int(!it.empty)
+
+function Base.iterate(it::OneOrZeroElementVector)
+    if it.empty
+        return nothing
+    else
+        return it.el, nothing
+    end
+end
+Base.iterate(::OneOrZeroElementVector, ::Nothing) = nothing
+
+Base.getindex(it::OneOrZeroElementVector, i::Integer) = it.el
+
+Iterators.reverse(it::OneOrZeroElementVector) = it
+
 # Copied from MathOptInterface
 """
     struct LazyMap{T, VT}

src/monovec.jl

@@ -1,6 +1,6 @@
 export monovec, monovectype, emptymonovec, sortmonovec, mergemonovec
 
-monomials(v::AbstractVariable, args...) = monomials((v,), args...)
+monomials(v::AbstractVariable, degree, args...) = monomials((v,), degree, args...)
 
 """
     emptymonovec(p::AbstractPolynomialLike)

src/polynomial.jl

@@ -147,7 +147,7 @@ Returns an iterator over the nonzero terms of the polynomial `p` sorted in the d
 
 Calling `terms` on ``4x^2y + xy + 2x`` should return an iterator of ``[4x^2y, xy, 2x]``.
 """
-terms(t::AbstractTermLike) = iszero(t) ? termtype(t)[] : [term(t)]
+terms(t::AbstractTermLike) = OneOrZeroElementVector(iszero(t), term(t))
 terms(p::AbstractPolynomialLike) = terms(polynomial(p))
 
 """
@@ -178,7 +178,7 @@ Returns an iterator over the coefficients of the monomials of `X` in `p` where `
 Calling `coefficients` on ``4x^2y + xy + 2x`` should return an iterator of ``[4, 1, 2]``.
 Calling `coefficients(4x^2*y + x*y + 2x + 3, [x, 1, x*y, y])` should return an iterator of ``[2, 3, 1, 0]``.
 """
-coefficients(p::APL) = coefficient.(terms(p))
+coefficients(p::APL{T}) where {T} = LazyMap{T}(coefficient, terms(p))
 function coefficients(p::APL{T}, X::AbstractVector) where T
     σ, mv = sortmonovec(X)
     @assert length(mv) == length(X) # no duplicate in X
@@ -214,6 +214,7 @@ Calling `monomials` on ``4x^2y + xy + 2x`` should return an iterator of ``[x^2y,
 Calling `monomials((x, y), [1, 3], m -> degree(m, y) != 1)` should return `[x^3, x*y^2, y^3, x]` where `x^2*y` and `y` have been excluded by the filter.
 """
 monomials(p::APL) = monovec(monomial.(terms(p)))
+monomials(t::AbstractTermLike) = OneOrZeroElementVector(iszero(t), monomial(t))
 
 function isconstant(p::APL)
     n = nterms(p)

src/term.jl

@@ -24,7 +24,7 @@ When applied to a monomial, it create a term of type `AbstractTerm{Int}`.
 function term end
 term(coef, var::AbstractVariable) = term(coef, monomial(var))
 term(coef, mono::AbstractMonomialLike) = termtype(mono, typeof(coef))(coef, mono)
-term(p::APL) = convert(termtype(p), p)
+term(p::APL) = convert(termtype(typeof(p)), p)
 
 """
     termtype(p::AbstractPolynomialLike)

test/allocations.jl

@@ -18,9 +18,13 @@ using TypedPolynomials
 function test_polynomial_merge()
     @polyvar x
     p = x^2 + x + 1
-    q = 2x + 3
-    alloc_test(0) do
-        add!!(p, q)
+    for q in [x, x^2, 2x, 2x + 3]
+        alloc_test(0) do
+            add!!(p, q)
+        end
+        alloc_test(0) do
+            sub!!(p, q)
+        end
     end
 end