Specialize on function type (#208)

* Specialize on function type

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

* Fix

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

Author: blegat

src/default_polynomial.jl

@@ -105,7 +105,7 @@ Base.:(-)(p1::Polynomial{<:LinearAlgebra.UniformScaling}, p2::Polynomial{<:Linea
 
 LinearAlgebra.adjoint(x::Polynomial) = polynomial!(adjoint.(terms(x)))
 
-function mapcoefficients(f::Function, p::Polynomial; nonzero = false)
+function mapcoefficients(f::F, p::Polynomial; nonzero = false) where {F<:Function}
     terms = map(p.terms) do term
         mapcoefficients(f, term)
     end
@@ -114,7 +114,7 @@ function mapcoefficients(f::Function, p::Polynomial; nonzero = false)
     end
     return polynomial!(terms)
 end
-function mapcoefficients!(f::Function, p::Polynomial; nonzero = false)
+function mapcoefficients!(f::F, p::Polynomial; nonzero = false) where {F<:Function}
     for i in eachindex(p.terms)
         t = p.terms[i]
         p.terms[i] = Term(f(coefficient(t)), monomial(t))
@@ -125,7 +125,7 @@ function mapcoefficients!(f::Function, p::Polynomial; nonzero = false)
     return p
 end
 
-function mapcoefficients_to!(output::Polynomial, f::Function, p::Polynomial; nonzero = false)
+function mapcoefficients_to!(output::Polynomial, f::F, p::Polynomial; nonzero = false) where {F<:Function}
     resize!(output.terms, nterms(p))
     for i in eachindex(p.terms)
         t = p.terms[i]

src/operators.jl

@@ -76,9 +76,9 @@ MA.operate_to!(output::AbstractPolynomial, op::typeof(*), p::APL, α) = MA.opera
 MA.operate_to!(output::APL, op::typeof(/), p::APL, α) = mapcoefficients_to!(output, Base.Fix2(op, α), p)
 
 function polynomial_merge!(
-    n1::Int, n2::Int, get1::Function, get2::Function,
-    set::Function, push::Function, compare_monomials::Function,
-    combine::Function, keep::Function, resize::Function)
+    n1::Int, n2::Int, get1::F1, get2::F2,
+    set::F3, push::F4, compare_monomials::F5,
+    combine::F6, keep::F7, resize::F8) where {F1, F2, F3, F4, F5, F6, F7, F8}
     buffer = nothing
     i = j = k = 1
     # Invariant:

src/polynomial.jl

@@ -58,7 +58,7 @@ function polynomial(Q::AbstractMatrix, mv::AbstractVector, ::Type{T}) where T
     polynomial(polynomial(Q, mv), T)
 end
 
-polynomial(f::Function, mv::AbstractVector{<:AbstractMonomialLike}) = polynomial!([term(f(i), mv[i]) for i in 1:length(mv)])
+polynomial(f::F, mv::AbstractVector{<:AbstractMonomialLike}) where {F<:Function} = polynomial!([term(f(i), mv[i]) for i in 1:length(mv)])
 
 function polynomial(a::AbstractVector, x::AbstractVector, s::ListState=MessyState())
     # If `x` is e.g. `[v, 1]` then it will contains terms that are convertible to monomials.
@@ -404,8 +404,8 @@ function _divtoone(t::AbstractTermLike{T}, α::S) where {T, S}
 end
 
 # TODO deprecate
-mapcoefficientsnz(f::Function, p::APL) = mapcoefficients(f, p, nonzero = true)
-mapcoefficientsnz_to!(output::APL, f::Function, p::APL) = mapcoefficients_to!(output, f, p, nonzero = true)
+mapcoefficientsnz(f::F, p::APL) where {F<:Function} = mapcoefficients(f, p, nonzero = true)
+mapcoefficientsnz_to!(output::APL, f::F, p::APL) where {F<:Function} = mapcoefficients_to!(output, f, p, nonzero = true)
 
 """
     mapcoefficients(f::Function, p::AbstractPolynomialLike, nonzero = false)
@@ -422,13 +422,13 @@ See also [`mapcoefficients!`](@ref) and [`mapcoefficients_to!`](@ref).
 Calling `mapcoefficients(α -> mod(3α, 6), 2x*y + 3x + 1)` should return `3x + 3`.
 """
 function mapcoefficients end
-function mapcoefficients(f::Function, p::AbstractPolynomialLike; nonzero = false) # Not used by either TypedPolynomials or DynamicPolynomials but used by CustomPoly in tests. FIXME Remove in a breaking release
+function mapcoefficients(f::F, p::AbstractPolynomialLike; nonzero = false) where {F<:Function} # Not used by either TypedPolynomials or DynamicPolynomials but used by CustomPoly in tests. FIXME Remove in a breaking release
     # Invariant: p has only nonzero coefficient
     # therefore f(α) will be nonzero for every coefficient α of p
     # hence we can use Uniq
     polynomial!(mapcoefficients.(f, terms(p)), nonzero ? SortedUniqState() : SortedState())
 end
-function mapcoefficients(f::Function, t::AbstractTermLike; nonzero = false)
+function mapcoefficients(f::F, t::AbstractTermLike; nonzero = false) where {F<:Function}
     return term(f(coefficient(t)), monomial(t))
 end