Improve type inference in mortar (#212)

* type inference in block sizes

* version bump to v0.16.17

* test for different-sized axes

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "BlockArrays"
 uuid = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
-version = "0.16.16"
+version = "0.16.17"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

src/blockarray.jl

@@ -266,8 +266,10 @@ function sizes_from_blocks(blocks::AbstractArray{<:Any, N}, _) where N
         error("All blocks must have ndims consistent with ndims = $N of `blocks` array.")
     end
     fullsizes = map!(size, Array{NTuple{N,Int}, N}(undef, size(blocks)), blocks)
-    block_sizes = ntuple(ndims(blocks)) do i
-        [s[i] for s in view(fullsizes, ntuple(j -> j == i ? (:) : 1, ndims(blocks))...)]
+    fR = reinterpret(reshape, Int, fullsizes)
+    stfR = strides(fR)
+    block_sizes = ntuple(N) do i
+        fR[range(i, step = stfR[i+1], length=size(fullsizes, i))]
     end
     checksizes(fullsizes, block_sizes)
     return block_sizes

test/test_blockarrays.jl

@@ -125,17 +125,22 @@ end
         @test_throws DimensionMismatch BlockArray([1,2,3],[1,1])
 
         @testset "mortar" begin
-            @testset for sizes in [(1:3,), (1:3, 1:3), (1:3, 1:3, 1:3)]
+            @testset for sizes in [(1:3,), (1:3, 1:4), (1:3, 1:4, 1:2)]
                 dims = sum.(sizes)
-                A = BlockArray(copy(reshape(1:prod(dims), dims)), sizes...)
-                @test mortar(A.blocks) == A
+                A = @inferred BlockArray(copy(reshape(1:prod(dims), dims)), sizes...)
+                @test @inferred mortar(A.blocks) == A
+                if length(dims) == 2
+                    # compare with hvcat
+                    rows = ntuple(_->length(sizes[2]), length(sizes[1]))
+                    @test mortar(A.blocks) == hvcat(rows, permutedims(A.blocks)...)
+                end
             end
 
-            ret = mortar([spzeros(2), spzeros(3)])
+            ret = @inferred mortar([spzeros(2), spzeros(3)])
             @test eltype(ret.blocks) <: SparseVector
             @test axes(ret) == (blockedrange([2, 3]),)
 
-            ret = mortar(
+            ret = @inferred mortar(
                 (spzeros(1, 3), spzeros(1, 4)),
                 (spzeros(2, 3), spzeros(2, 4)),
                 (spzeros(5, 3), spzeros(5, 4)),
@@ -157,6 +162,17 @@ end
                     (zeros(2, 3), zeros(111, 222)),
                 )
             end
+
+            @testset "sizes_from_blocks" begin
+                blocks = reshape([rand(2,2), zeros(1,2),
+                                  zeros(2,3), rand(1,3)], 2, 2);
+                @test @inferred BlockArrays.sizes_from_blocks(blocks) == ([2,1], [2,3])
+                blocks = reshape(
+                    [rand(2,2), zeros(1,2), zeros(4,2),
+                     zeros(2,3), rand(1,3), zeros(4,3),
+                     zeros(2,1), zeros(1,1), rand(4,1)], 3, 3);
+                @test @inferred BlockArrays.sizes_from_blocks(blocks) == ([2, 1, 4], [2, 3, 1])
+            end
         end
 
         @testset "BlockVector" begin