Fix iteration for SparsePolynomial

Project.toml

@@ -2,7 +2,7 @@ name = "Polynomials"
 uuid = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
 license = "MIT"
 author = "JuliaMath"
-version = "3.0.1"
+version = "3.1.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/common.jl

@@ -460,8 +460,8 @@ Base.length(p::AbstractPolynomial) = length(coeffs(p))
 
 Returns the size of the polynomials coefficients, along axis `i` if provided.
 """
-Base.size(p::AbstractPolynomial) = size(coeffs(p))
-Base.size(p::AbstractPolynomial, i::Integer) = size(coeffs(p), i)
+Base.size(p::AbstractPolynomial) = (length(p),)
+Base.size(p::AbstractPolynomial, i::Integer) =  i <= 1 ? size(p)[i] : 1
 Base.eltype(p::AbstractPolynomial{T}) where {T} = T
 # in  analogy  with  polynomial as a Vector{T} with different operations defined.
 Base.eltype(::Type{<:AbstractPolynomial}) = Float64
@@ -470,7 +470,7 @@ _eltype(::Type{<:AbstractPolynomial}) = nothing
 _eltype(::Type{<:AbstractPolynomial{T}}) where {T} = T
 function _eltype(P::Type{<:AbstractPolynomial}, p::AbstractPolynomial)
     T′ = _eltype(P)
-    T = T′ == nothing ? eltype(p) : T′
+    T = T′ === nothing ? eltype(p) : T′
     T
 end
 Base.iszero(p::AbstractPolynomial) = all(iszero, p)
@@ -664,47 +664,34 @@ Iteration =#
 # `pairs(p)`: `i => pᵢ` possibly skipping over values of `i` with `pᵢ == 0` (SparsePolynomial)
 #    and possibly non ordered (SparsePolynomial)
 # `monomials(p)`: iterates over pᵢ ⋅ basis(p, i) i  ∈ keys(p)
-function Base.iterate(p::AbstractPolynomial, state=nothing)
-    i = firstindex(p)
-    if state == nothing
-        return (p[i], i)
-    else
-        j = lastindex(p)
-        if i <= state < j
-            return (p[state+1], state+1)
-        end
-        return nothing
-    end
+function _iterate(p, state)
+    firstindex(p) <= state <= lastindex(p) || return nothing
+    return p[state], state+1
 end
+Base.iterate(p::AbstractPolynomial, state = firstindex(p)) = _iterate(p, state)
 
 # pairs map i -> aᵢ *possibly* skipping over ai == 0
 # cf. abstractdict.jl
-struct PolynomialKeys{P}
+struct PolynomialKeys{P} <: AbstractSet{Int}
     p::P
 end
-struct PolynomialValues{P}
+struct PolynomialValues{P, T} <: AbstractSet{T}
     p::P
+
+    PolynomialValues{P}(p::P) where {P} = new{P, eltype(p)}(p)
+    PolynomialValues(p::P) where {P} = new{P, eltype(p)}(p)
 end
 Base.keys(p::AbstractPolynomial) =  PolynomialKeys(p)
 Base.values(p::AbstractPolynomial) =  PolynomialValues(p)
 Base.length(p::PolynomialValues) = length(p.p.coeffs)
 Base.length(p::PolynomialKeys) = length(p.p.coeffs)
 Base.size(p::Union{PolynomialValues, PolynomialKeys}) = (length(p),)
-function Base.iterate(v::PolynomialKeys, state=nothing)
-    i = firstindex(v.p)
-    state==nothing && return (i, i)
-    j = lastindex(v.p)
-    i <= state < j && return (state+1, state+1)
-    return nothing
+function Base.iterate(v::PolynomialKeys, state = firstindex(v.p))
+    firstindex(v.p) <= state <= lastindex(v.p) || return nothing
+    return state, state+1
 end
 
-function Base.iterate(v::PolynomialValues, state=nothing)
-    i = firstindex(v.p)
-    state==nothing && return (v.p[i], i)
-    j = lastindex(v.p)
-    i <= state < j && return (v.p[state+1], state+1)
-    return nothing
-end
+Base.iterate(v::PolynomialValues, state = firstindex(v.p)) = _iterate(v.p, state)
 
 
 # iterate over monomials of the polynomial
@@ -719,7 +706,7 @@ Returns an iterator over the terms, `pᵢ⋅basis(p,i)`, of the polynomial for e
 monomials(p) = Monomials(p)
 function Base.iterate(v::Monomials, state...)
     y = iterate(pairs(v.p), state...)
-    y == nothing && return nothing
+    y === nothing && return nothing
     kv, s = y
     return (kv[2]*basis(v.p, kv[1]), s)
 end
@@ -736,18 +723,18 @@ _indeterminate(::Type{P}) where {P <: AbstractPolynomial} = nothing
 _indeterminate(::Type{P}) where {T, X, P <: AbstractPolynomial{T,X}} = X
 function indeterminate(::Type{P}) where {P <: AbstractPolynomial}
     X = _indeterminate(P)
-    X == nothing ? :x : X
+    X === nothing ? :x : X
 end
 indeterminate(p::P) where {P <: AbstractPolynomial} = _indeterminate(P)
 function indeterminate(PP::Type{P}, p::AbstractPolynomial{T,Y}) where {P <: AbstractPolynomial, T,Y}
     X = _indeterminate(PP)
-    X == nothing && return Y
+    X === nothing && return Y
     assert_same_variable(X,Y)
     return X
-    #X = _indeterminate(PP) == nothing ? indeterminate(p) :  _indeterminate(PP)
+    #X = _indeterminate(PP) === nothing ? indeterminate(p) :  _indeterminate(PP)
 end
 function indeterminate(PP::Type{P}, x::Symbol) where {P <: AbstractPolynomial}
-    X = _indeterminate(PP) == nothing ? x :  _indeterminate(PP)
+    X = _indeterminate(PP) === nothing ? x :  _indeterminate(PP)
 end
 
 #=
@@ -774,7 +761,7 @@ Base.zero(p::P, var=indeterminate(p)) where {P <: AbstractPolynomial} = zero(P,
 Returns a representation of 1 as the given polynomial.
 """
 Base.one(::Type{P}) where {P<:AbstractPolynomial} = throw(ArgumentError("No default method defined")) # no default method
-Base.one(::Type{P}, var::SymbolLike) where {P <: AbstractPolynomial} = one(⟒(P){eltype(P), Symbol(var == nothing ? :x : var)})
+Base.one(::Type{P}, var::SymbolLike) where {P <: AbstractPolynomial} = one(⟒(P){eltype(P), Symbol(var === nothing ? :x : var)})
 Base.one(p::P, var=indeterminate(p)) where {P <: AbstractPolynomial} = one(P, var)
 
 Base.oneunit(::Type{P}, args...) where {P <: AbstractPolynomial} = one(P, args...)

src/polynomials/Poly.jl

@@ -47,13 +47,9 @@ _eltype(::Type{<:Poly{T}}) where  {T} = T
 _eltype(::Type{Poly}) =  Float64
 
 # when interating over poly return monomials
-function Base.iterate(p::Poly, state=nothing)
-    i = 0
-    state == nothing && return (p[i]*one(p), i)
-    j = degree(p)
-    s = state + 1
-    i <= state < j && return (p[s]*Polynomials.basis(p,s), s)
-    return nothing
+function Base.iterate(p::Poly, state = firstindex(p))
+    firstindex(p) <= state <= lastindex(p) || return nothing
+    return p[state] * Polynomials.basis(p,state), state+1
 end
 Base.collect(p::Poly) = [pᵢ for pᵢ ∈ p]
 

src/polynomials/SparsePolynomial.jl

@@ -127,11 +127,6 @@ function Base.setindex!(p::SparsePolynomial, value::Number, idx::Int)
     return p
 end
 
-
-Base.firstindex(p::SparsePolynomial) = sort(collect(keys(p.coeffs)), by=x->x[1])[1]
-Base.lastindex(p::SparsePolynomial) = sort(collect(keys(p.coeffs)), by=x->x[1])[end]
-Base.eachindex(p::SparsePolynomial) = sort(collect(keys(p.coeffs)), by=x->x[1])
-
 # pairs iterates only over non-zero
 # inherits order for underlying dictionary
 function Base.iterate(v::PolynomialKeys{SparsePolynomial{T,X}}, state...) where {T,X}
@@ -146,7 +141,7 @@ function Base.iterate(v::PolynomialValues{SparsePolynomial{T,X}}, state...) wher
     return (y[1][2], y[2])
 end
 
-
+Base.length(S::SparsePolynomial) = isempty(S.coeffs) ? 0 : maximum(keys(S.coeffs)) + 1
 
 ##
 ## ----

test/ChebyshevT.jl

@@ -1,19 +1,22 @@
-@testset "Construction" for coeff in [
-    Int64[1, 1, 1, 1],
-    Float32[1, -4, 2],
-    ComplexF64[1 - 1im, 2 + 3im],
-    [3 // 4, -2 // 1, 1 // 1]
-]
-    p = ChebyshevT(coeff)
-    @test p.coeffs == coeff
-    @test coeffs(p) == coeff
-    @test degree(p) == length(coeff) - 1
-    @test Polynomials.indeterminate(p) == :x
-    @test length(p) == length(coeff)
-    @test size(p) == size(coeff)
-    @test size(p, 1) == size(coeff, 1)
-    @test typeof(p).parameters[1] == eltype(coeff)
-    @test eltype(p) == eltype(coeff)
+@testset "Construction" begin
+    @testset for coeff in Any[
+        Int64[1, 1, 1, 1],
+        Float32[1, -4, 2],
+        ComplexF64[1 - 1im, 2 + 3im],
+        [3 // 4, -2 // 1, 1 // 1]
+    ]
+
+        p = ChebyshevT(coeff)
+        @test p.coeffs == coeff
+        @test coeffs(p) == coeff
+        @test degree(p) == length(coeff) - 1
+        @test Polynomials.indeterminate(p) == :x
+        @test length(p) == length(coeff)
+        @test size(p) == size(coeff)
+        @test size(p, 1) == size(coeff, 1)
+        @test typeof(p).parameters[1] == eltype(coeff)
+        @test eltype(p) == eltype(coeff)
+    end
 end
 
 @testset "Other Construction" begin

test/StandardBasis.jl

@@ -28,36 +28,39 @@ Ps = (ImmutablePolynomial, Polynomial, SparsePolynomial, LaurentPolynomial, Fact
 isimmutable(p::P) where {P} = P <: ImmutablePolynomial
 isimmutable(::Type{<:ImmutablePolynomial}) = true
 
-@testset "Construction" for coeff in [
-                                      Int64[1, 1, 1, 1],
-                                      Float32[1, -4, 2],
-                                      ComplexF64[1 - 1im, 2 + 3im],
-                                      [3 // 4, -2 // 1, 1 // 1]
-                                     ]
+@testset "Construction" begin
+    @testset for coeff in Any[
+          Int64[1, 1, 1, 1],
+          Float32[1, -4, 2],
+          ComplexF64[1 - 1im, 2 + 3im],
+          [3 // 4, -2 // 1, 1 // 1]
+         ]
+
+        @testset for P in Ps
+            p = P(coeff)
+            @test coeffs(p) ==ᵗ⁰ coeff
+            @test degree(p) == length(coeff) - 1
+            @test indeterminate(p) == :x
+            P == Polynomial && @test length(p) == length(coeff)
+            P == Polynomial && @test size(p) == size(coeff)
+            P == Polynomial && @test size(p, 1) == size(coeff, 1)
+            P == Polynomial && @test typeof(p).parameters[1] == eltype(coeff)
+            if !(eltype(coeff) <: Real && P == FactoredPolynomial) # roots may be complex
+                @test eltype(p) == eltype(coeff)
+            end
+            @test all([-200, -0.3, 1, 48.2] .∈ Polynomials.domain(p))
 
-    for P in Ps
-        p = P(coeff)
-        @test coeffs(p) ==ᵗ⁰ coeff
-        @test degree(p) == length(coeff) - 1
-        @test indeterminate(p) == :x
-        P == Polynomial && @test length(p) == length(coeff)
-        P == Polynomial && @test size(p) == size(coeff)
-        P == Polynomial && @test size(p, 1) == size(coeff, 1)
-        P == Polynomial && @test typeof(p).parameters[1] == eltype(coeff)
-        @test eltype(p) == eltype(coeff)
-        @test all([-200, -0.3, 1, 48.2] .∈ Polynomials.domain(p))
-
-        ## issue #316
-        @test_throws InexactError P{Int,:x}([1+im, 1])
-        @test_throws InexactError P{Int}([1+im, 1], :x)
-        @test_throws InexactError P{Int,:x}(1+im)
-        @test_throws InexactError P{Int}(1+im)
-
-        ## issue #395
-        v = [1,2,3]
-        @test P(v) == P(v,:x) == P(v,'x') == P(v,"x") == P(v, Polynomials.Var(:x))
-    end
+            ## issue #316
+            @test_throws InexactError P{Int,:x}([1+im, 1])
+            @test_throws InexactError P{Int}([1+im, 1], :x)
+            @test_throws InexactError P{Int,:x}(1+im)
+            @test_throws InexactError P{Int}(1+im)
 
+            ## issue #395
+            v = [1,2,3]
+            @test P(v) == P(v,:x) == P(v,'x') == P(v,"x") == P(v, Polynomials.Var(:x))
+        end
+    end
 end
 
 @testset "Mapdomain" begin
@@ -1493,3 +1496,13 @@ end
     @test c == MA.operate(*, A, b)
     @test 0 == @allocated MA.buffered_operate!(buffer, MA.add_mul, c, A, b)
 end
+
+@testset "empty SparsePolynomial" begin
+    p = SparsePolynomial(Float64[0])
+    @test eltype(p) == Float64
+    @test eltype(keys(p)) == Int
+    @test eltype(values(p)) == Float64
+    @test collect(p) == Float64[]
+    @test collect(keys(p)) == Int[]
+    @test collect(values(p)) == Float64[]
+end