Add blockcheckindex and blockcheckbounds_indices

Author: jishnub

src/abstractblockarray.jl

@@ -87,32 +87,77 @@ ERROR: BlockBoundsError: attempt to access 2×2-blocked 2×3 BlockMatrix{Float64
 [...]
 ```
 """
-@inline function blockcheckbounds(A::AbstractArray{T, N}, i::Vararg{Integer, N}) where {T,N}
-    if blockcheckbounds(Bool, A, i...)
-        return
-    else
-        throw(BlockBoundsError(A, i))
-    end
+@inline function blockcheckbounds(A::AbstractArray, i::Integer...)
+    blockcheckbounds(Bool, A, i...) || throw(BlockBoundsError(A, i))
 end
 
-@inline function blockcheckbounds(::Type{Bool}, A::AbstractArray{T, N}, i::Vararg{Integer, N}) where {T,N}
-    n = blockaxes(A)
-    k = 0
-    for idx in 1:N # using enumerate here will allocate
-        k += 1
-        @inbounds _i = i[idx]
-        Block(_i) in n[k] || return false
-    end
-    for idx = N+1:length(i)
-        isone(i[idx]) || return false
-    end
-    return true
+# linear block indexing
+@inline function blockcheckbounds(::Type{Bool}, A::AbstractArray, i::Integer)
+    blockcheckindex(Bool, BlockRange(blocklength(A)), i)
+end
+# cartesian block indexing
+@inline function blockcheckbounds(::Type{Bool}, A::AbstractArray, i::Integer...)
+    blockcheckbounds_indices(Bool, blockaxes(A), i)
 end
 
 blockcheckbounds(A::AbstractArray{T, N}, i::Block{N}) where {T,N} = blockcheckbounds(A, i.n...)
 blockcheckbounds(A::AbstractArray{T, N}, i::Vararg{Block{1},N}) where {T,N} = blockcheckbounds(A, Int.(i)...)
 blockcheckbounds(A::AbstractVector{T}, i::Block{1}) where {T} = blockcheckbounds(A, Int(i))
 
+"""
+    blockcheckbounds_indices(Bool, IA::Tuple{Vararg{BlockRange{1}}}, I::Tuple{Vararg{Integer}})
+
+Return true if the "requested" indices in the tuple `Block.(I)` fall within the bounds of the "permitted"
+indices specified by the tuple `IA`. This function recursively consumes elements of these tuples
+in a 1-for-1 fashion.
+
+The actual bounds-checking is performed by [`blockcheckindex`](@ref).
+
+# Examples
+```jldoctest
+julia> B = BlockArray(zeros(6,6), 1:3, 1:3);
+
+julia> blockaxes(B)
+(BlockRange(Base.OneTo(3)), BlockRange(Base.OneTo(3)))
+
+julia> BlockArrays.blockcheckbounds_indices(Bool, blockaxes(B), (1,2))
+true
+
+julia> BlockArrays.blockcheckbounds_indices(Bool, blockaxes(B), (4,1))
+false
+```
+"""
+@inline blockcheckbounds_indices(::Type{Bool}, ::Tuple{}, ::Tuple{}) = true
+@inline function blockcheckbounds_indices(::Type{Bool}, blockaxes::Tuple{Vararg{BlockRange{1}}}, i::Tuple{})
+    # the trailing blocks must be Block(1)
+    b = first(blockaxes)
+    length(b) == 1 && Int(b[]) == 1 && blockcheckbounds_indices(Bool, blockaxes[2:end], i)
+end
+@inline function blockcheckbounds_indices(::Type{Bool}, blockaxes::Tuple{}, i::Tuple{Vararg{Integer}})
+    # the trailing indices must be 1
+    first(i) == 1 && blockcheckbounds_indices(Bool, blockaxes, i[2:end])
+end
+@inline function blockcheckbounds_indices(::Type{Bool}, blockaxes::Tuple{Vararg{BlockRange{1}}}, i::Tuple{Vararg{Integer}})
+    blockcheckindex(Bool, first(blockaxes), first(i)) &&
+        blockcheckbounds_indices(Bool, blockaxes[2:end], i[2:end])
+end
+
+"""
+    blockcheckindex(Bool, inds::BlockRange{1}, index::Integer)
+
+Return `true` if `Block(index)` is within the bounds of `inds`.
+
+# Examples
+```jldoctest
+julia> BlockArrays.blockcheckindex(Bool, BlockRange(1:2), 1)
+true
+
+julia> BlockArrays.blockcheckindex(Bool, BlockRange(1:2), 3)
+false
+```
+"""
+@inline blockcheckindex(::Type{Bool}, inds::BlockRange{1}, i::Integer) = Block(i) in inds
+
 @propagate_inbounds Base.setindex!(block_arr::AbstractBlockArray{T,N}, v, block::Block{N}) where {T,N} =
     setindex!(block_arr, v, Block.(block.n)...)
 @inline @propagate_inbounds function Base.setindex!(block_arr::AbstractBlockArray{T,N}, v, block::Vararg{Block{1}, N}) where {T,N}

test/test_blockarrays.jl

@@ -277,6 +277,13 @@ end
         @test_throws DimensionMismatch (BA_1[Block(3)] = rand(4))
         @test_throws BlockBoundsError blockcheckbounds(BA_1, 4)
         @test_throws BlockBoundsError BA_1[Block(4)]
+        @test_throws BlockBoundsError blockcheckbounds(BA_1)
+        @test_throws BlockBoundsError blockcheckbounds(BA_1, 4, 1)
+
+        Bv = BlockArray(zeros(1))
+        @test Bv[Block()] == Bv[] == 0
+        @test Bv[Block(1)] == Bv
+        @test Bv[Block(1,1)] == zeros(1,1)
 
         BA_2 = BlockArray(undef_blocks, Matrix{Float64}, [1,2], [3,4])
         a_2 = rand(1,4)
@@ -286,6 +293,26 @@ end
 
         @test BA_2[1,5] == a_2[2]
         @test_throws DimensionMismatch BA_2[Block(1,2)] = rand(1,5)
+
+        # linear block indexing
+        @test blockcheckbounds(Bool, BA_2, 3)
+        @test_throws BlockBoundsError blockcheckbounds(BA_2, 5)
+
+        # trailing Block(1) indices
+        BA_3 = BlockArray(undef_blocks, Matrix{Float64}, [1,3], [4])
+        @test blockcheckbounds(Bool, BA_3, 1, 1)
+        @test blockcheckbounds(Bool, BA_3, 2, 1)
+        @test_throws BlockBoundsError blockcheckbounds(BA_3, 3, 1)
+        @test_throws BlockBoundsError blockcheckbounds(BA_3, 1, 2)
+        @test_throws BlockBoundsError blockcheckbounds(BA_3, 3, 2)
+
+        BA_4 = BlockArray(zeros(2,2,1))
+        @test blockcheckbounds(Bool, BA_4)
+        @test blockcheckbounds(Bool, BA_4, 1)
+        @test blockcheckbounds(Bool, BA_4, 1, 1)
+        @test blockcheckbounds(Bool, BA_4, 1, 1, 1)
+        @test blockcheckbounds(Bool, BA_4, 1, 1, 1, 1)
+        @test_throws BlockBoundsError blockcheckbounds(BA_3, 1, 2)
     end
 
     @testset "misc block tests" begin