improve TensorSpace inference (#261)

Author: jishnub

Project.toml

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

src/Multivariate/TensorSpace.jl

@@ -222,11 +222,13 @@ tensor_eval_type(::Type{Vector{Any}},::Type{Vector{Any}}) = Vector{Any}
 tensor_eval_type(::Type{Vector{Any}},_) = Vector{Any}
 tensor_eval_type(_,::Type{Vector{Any}}) = Vector{Any}
 
-
+# Specialize some common cases to avoid mapreduce, which has inference issues
+_typeofproddomain(sp::Tuple{Any}) = typeof(domain(sp[1]))
+_typeofproddomain(sp::Tuple{Any,Any}) = typeof(domain(sp[1]) × domain(sp[2]))
+_typeofproddomain(sp) = typeof(mapreduce(domain,×,sp))
 TensorSpace(sp::Tuple) =
-    TensorSpace{typeof(sp),typeof(mapreduce(domain,×,sp)),
-                mapreduce(rangetype,(a,b)->tensor_eval_type(a,b),sp)}(sp)
-
+    TensorSpace{typeof(sp), _typeofproddomain(sp),
+                mapreduce(rangetype,tensor_eval_type,sp)}(sp)
 
 dimension(sp::TensorSpace) = mapreduce(dimension,*,sp.spaces)
 

test/SpacesTest.jl

@@ -253,6 +253,12 @@ using LinearAlgebra
         @test dimension(a^3) == dimension(a)^3
         @test @inferred(domain(a^3)) == domain(a)^3
         @test_broken @inferred(points(a^3)) == vec(Vec.(points(a), points(a)', reshape(points(a), 1,1,4)))
+
+        p = PointSpace(1:4)
+        d = domain(p)
+        @test domain(TensorSpace(p)) == d
+        @test components(domain(TensorSpace(p, p))) == (d, d)
+        @test components(domain(TensorSpace(p, p, p))) == (d, d, d)
     end
 
     @testset "ConstantSpace" begin