From eb95d65b5317f94898cdc9156f4591fcf492516d Mon Sep 17 00:00:00 2001
From: Jishnu Bhattacharya <jishnub.github@gmail.com>
Date: Thu, 7 Nov 2024 12:12:14 +0530
Subject: [PATCH] Inline blocksizes for better type-inference

---
 src/blocks.jl       | 2 +-
 test/test_blocks.jl | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/blocks.jl b/src/blocks.jl
index 711a182a..b87550ca 100644
--- a/src/blocks.jl
+++ b/src/blocks.jl
@@ -125,7 +125,7 @@ julia> blocksizes(A,2)
 ```
 """
 blocksizes(A::AbstractArray) = BlockSizes(A)
-blocksizes(A::AbstractArray, d::Integer) = blocklengths(axes(A, d))
+@inline blocksizes(A::AbstractArray, d::Integer) = blocklengths(axes(A, d))
 
 struct BlockSizes{T,N,A<:AbstractArray{<:Any,N}} <: AbstractArray{T,N}
     array::A
diff --git a/test/test_blocks.jl b/test/test_blocks.jl
index 9a636bf9..3ad25e64 100644
--- a/test/test_blocks.jl
+++ b/test/test_blocks.jl
@@ -127,6 +127,14 @@ end
         @test blocksizes(A, 1) == [2, 3]
         @test blocksizes(A, 2) == [3, 1]
     end
+
+    @testset "Inference: issue #425" begin
+        x = BlockedArray(rand(4, 4), [2, 2], [2, 2])
+        bs1 = @inferred (x -> blocksizes(x, 1))(x)
+        @test bs1 == [2,2]
+        bs4 = @inferred (x -> blocksizes(x, 4))(x)
+        @test bs4 == 1:1
+    end
 end
 
 end # module