Improve inference in PlusOperator (#269)

* Tuples in PlusOperator

* Add inference tests

* Generalize collateop

* Version bump to v0.7.36

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "ApproxFunBase"
 uuid = "fbd15aa5-315a-5a7d-a8a4-24992e37be05"
-version = "0.7.35"
+version = "0.7.36"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/Operators/general/algebra.jl

@@ -46,7 +46,8 @@ function convert(::Type{Operator{T}}, P::PlusOperator) where T
         P
     else
         ops = P.ops
-        PlusOperator(ops isa AbstractVector{<:Operator{T}} ? ops : map(x -> strictconvert(Operator{T}, x), ops),
+        PlusOperator(ops isa AbstractVector{<:Operator{T}} ? ops :
+                map(x -> strictconvert(Operator{T}, x), ops),
             P.bandwidths,P.sz)::Operator{T}
     end
 end
@@ -64,19 +65,19 @@ end
 domain(P::PlusOperator) = commondomain(P.ops)
 
 _promote_eltypeof(As...) = _promote_eltypeof(As)
-_promote_eltypeof(As::Union{Vector, Tuple}) = mapreduce(eltype, promote_type, As)
-_promote_eltypeof(As::Vector{Operator{T}}) where {T} = T
+_promote_eltypeof(As::Union{AbstractVector, Tuple}) = mapreduce(eltype, promote_type, As)
+_promote_eltypeof(As::AbstractVector{<:Operator{T}}) where {T} = T
 
-_extractops(A, ::Any) = [A]
+_extractops(A, ::Any) = SVector{1}(A)
 _extractops(A::PlusOperator, ::typeof(+)) = A.ops
 
 function +(A::Operator,B::Operator)
-    v = [_extractops(A, +); _extractops(B, +)]
+    v = collateops(+, A, B)
     promoteplus(v, size(A))
 end
 # Optimization for 3-term sum
 function +(A::Operator,B::Operator,C::Operator)
-    v = [_extractops(A,+); _extractops(B, +); _extractops(C, +)]
+    v = collateops(+, A, B, C)
     promoteplus(v, size(A))
 end
 
@@ -247,10 +248,10 @@ __bandwidthssum(A, B::NTuple{2,InfiniteCardinal{0}}) = B
 __bandwidthssum(A, B) = reduce((t1, t2) -> t1 .+ t2, (A, B), init = (0,0))
 
 _timessize(ops) = (size(first(ops),1), size(last(ops),2))
-function TimesOperator(ops::Vector{O},
+function TimesOperator(ops::AbstractVector{O},
         bi::Tuple{Any,Any} = bandwidthssum(ops),
         sz::Tuple{Any,Any} = _timessize(ops)) where {T,O<:Operator{T}}
-    TimesOperator{T,typeof(bi),typeof(sz),O}(ops,bi,sz)
+    TimesOperator{T,typeof(bi),typeof(sz),O}(convert_vector(ops),bi,sz)
 end
 
 _extractops(A::TimesOperator, ::typeof(*)) = A.ops
@@ -268,7 +269,8 @@ function convert(::Type{Operator{T}},P::TimesOperator) where T
         P
     else
         ops = P.ops
-        TimesOperator(ops isa AbstractVector{<:Operator{T}} ? ops : map(x -> strictconvert(Operator{T}, x), ops) ,
+        TimesOperator(ops isa AbstractVector{<:Operator{T}} ? ops :
+                map(x -> strictconvert(Operator{T}, x), ops) ,
             bandwidths(P), size(P))
     end
 end
@@ -517,17 +519,22 @@ for OP in (:(adjoint),:(transpose))
         reverse(bandwidths(A)), reverse(size(A)))
 end
 
-_collateops(A::TimesOperator, B::TimesOperator, ::typeof(*)) = [_extractops(A, *); _extractops(B, *)]
-_collateops(A::TimesOperator, B::Operator, ::typeof(*)) = [_extractops(A, *); _extractops(B, *)]
-_collateops(A::Operator, B::TimesOperator, ::typeof(*)) = [_extractops(A, *); _extractops(B, *)]
-_collateops(A::Operator, B::Operator, ::typeof(*)) = (A, B)
+const PlusOrTimesOp = Union{PlusOperator, TimesOperator}
+anyplustimes(op::Operator, ops...) = anyplustimes(ops...)
+anyplustimes(op::PlusOrTimesOp, ops...) = true
+anyplustimes() = false
+
+collateops(op, As...) = collateops(op, Val(anyplustimes(As...)), As...)
+collateops(op, ::Val{true}, As...) = mapreduce(x -> _extractops(x, op), vcat, As)
+collateops(op, ::Val{false}, As...) = As
+
 function *(A::Operator,B::Operator)
     if isconstop(A)
         promoterangespace(strictconvert(Number,A)*B,rangespace(A))
     elseif isconstop(B)
         promotedomainspace(strictconvert(Number,B)*A,domainspace(B))
     else
-        promotetimes(_collateops(A, B, *),
+        promotetimes(collateops(*, A, B),
             domainspace(B), _timessize((A,B)), false)
     end
 end

test/runtests.jl

@@ -213,7 +213,7 @@ end
             f = Fun(PointSpace(1:3), c)
             M = Multiplication(f)
             @testset for t in [1, 3]
-                op = M + t * M
+                op = @inferred M + t * M
                 @test bandwidths(op) == bandwidths(M)
                 @test coefficients(op * f) == @. (1+t)*c^2
                 for op2 in Any[M + M + t * M, op + M]
@@ -230,9 +230,11 @@ end
                 @test coefficients(f1) == coefficients(f2) == coefficients(f3)
             end
             Z = ApproxFunBase.ZeroOperator()
-            @test Z + Z == Z
-            @test Z + Z + Z == Z
-            @test Z + Z + Z + Z == Z
+            @test (@inferred Z + Z) == Z
+            @test (@inferred Z + Z + Z) == Z
+            @test (@inferred Z + Z + Z + Z) == Z
+
+            @inferred (() -> (D = Derivative(); D + D))()
         end
     end