More general checkbounds and blockcheckbounds (#469)

This adds support for more bounds checking, such as `checkbounds(A,
Block.(2:3), Block.(2:3))`, `checkbounds(A, Block(2, 2)[1:3, 1:3])`,
etc.

It also fixes #458.

---------

Co-authored-by: Sheehan Olver <[email protected]>

Author: mtfishman

src/abstractblockarray.jl

@@ -88,6 +88,7 @@ ERROR: BlockBoundsError: attempt to access 2×2-blocked 2×3 BlockMatrix{Float64
 """
 @inline function blockcheckbounds(A::AbstractArray, i...)
     blockcheckbounds(Bool, A, i...) || throw(BlockBoundsError(A, i))
+    nothing
 end
 
 # linear block indexing
@@ -99,15 +100,29 @@ end
     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(::AbstractArray{T, 0}) where {T} = true
-blockcheckbounds(A::AbstractVector{T}, i::Block{1}) where {T} = blockcheckbounds(A, Int(i))
+# Used to ensure a `BlockBoundsError` is thrown instead of a `BoundsError`,
+# see https://github.com/JuliaArrays/BlockArrays.jl/issues/458
+checkbounds(A::AbstractArray{T, N}, I::Block{N}) where {T,N} = blockcheckbounds(A, I)
+checkbounds(A::AbstractArray{T}, I1::Block{1}, Irest::Block{1}...) where {T} = blockcheckbounds(A, I1, Irest...)
+checkbounds(A::AbstractArray{T}, I1::AbstractVector{<:Block{1}}, Irest::AbstractVector{<:Block{1}}...) where {T} =
+    blockcheckbounds(A, I1, Irest...)
+
+blockcheckbounds(::Type{Bool}, A::AbstractArray{T, N}, I::Block{N}) where {T,N} = blockcheckbounds(Bool, A, Int.(Tuple(I))...)
+blockcheckbounds(::Type{Bool}, A::AbstractArray{T, N}, I::Vararg{Block{1},N}) where {T,N} =
+    blockcheckbounds(Bool, A, Int.(I)...)
+blockcheckbounds(::Type{Bool}, ::AbstractArray{T, 0}) where {T} = true
+blockcheckbounds(::Type{Bool}, A::AbstractVector{T}, I::Block{1}) where {T} = blockcheckbounds(Bool, A, Int(I))
+blockcheckbounds(::Type{Bool}, A::AbstractArray{T,N}, I::Vararg{AbstractVector{<:Block{1}},N}) where {T,N} =
+    blockcheckbounds(Bool, A, map(i -> Int.(i), I)...)
+
+blockcheckbounds(::Type{Bool}, A::AbstractArray{T,N}, I::Vararg{BlockRange{1},N}) where {T,N} =
+    blockcheckbounds(Bool, A, map(i -> Int.(i), I)...)
+blockcheckbounds(::Type{Bool}, A::AbstractArray{T,N}, I::BlockRange{N}) where {T,N} = blockcheckbounds(Bool, A, I.indices...)
 
 """
-    blockcheckbounds_indices(Bool, IA::Tuple{Vararg{BlockRange{1}}}, I::Tuple{Vararg{Integer}})
+    blockcheckbounds_indices(Bool, IA::Tuple{Vararg{BlockRange{1}}}, I::Tuple)
 
-Return true if the "requested" indices in the tuple `Block.(I)` fall within the bounds of the "permitted"
+Return true if the "requested" indices in the tuple `map(i -> Block.(i), 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.
 
@@ -125,6 +140,12 @@ true
 
 julia> BlockArrays.blockcheckbounds_indices(Bool, blockaxes(B), (4,1))
 false
+
+julia> BlockArrays.blockcheckbounds_indices(Bool, blockaxes(B), (1:2,2:3))
+true
+
+julia> BlockArrays.blockcheckbounds_indices(Bool, blockaxes(B), (1:2,2:4))
+false
 ```
 """
 @inline blockcheckbounds_indices(::Type{Bool}, ::Tuple{}, ::Tuple{}) = true
@@ -143,9 +164,9 @@ end
 end
 
 """
-    blockcheckindex(Bool, inds::BlockRange{1}, index::Integer)
+    blockcheckindex(Bool, inds::BlockRange{1}, index)
 
-Return `true` if `Block(index)` is within the bounds of `inds`.
+Return `true` if `Block.(index)` is within the bounds of `inds`.
 
 # Examples
 ```jldoctest
@@ -154,9 +175,15 @@ true
 
 julia> BlockArrays.blockcheckindex(Bool, BlockRange(1:2), 3)
 false
+
+julia> BlockArrays.blockcheckindex(Bool, BlockRange(1:3), 2:3)
+true
+
+julia> BlockArrays.blockcheckindex(Bool, BlockRange(1:3), 2:4)
+false
 ```
 """
-@inline blockcheckindex(::Type{Bool}, inds::BlockRange{1}, i::Integer) = Block(i) in inds
+@inline blockcheckindex(::Type{Bool}, inds::BlockRange{1}, i) = checkindex(Bool, Int.(inds), i)
 
 @propagate_inbounds setindex!(block_arr::AbstractBlockArray{T,N}, v, block::Block{N}) where {T,N} =
     setindex!(block_arr, v, Block.(block.n)...)

src/blockindices.jl

@@ -164,6 +164,7 @@ end
 
 block(b::BlockIndex) = Block(b.I...)
 blockindex(b::BlockIndex{1}) = b.α[1]
+blockindex(b::BlockIndex) = CartesianIndex(b.α)
 
 BlockIndex(indcs::Tuple{Vararg{BlockIndex{1},N}}) where N = BlockIndex(block.(indcs), blockindex.(indcs))
 
@@ -172,15 +173,16 @@ BlockIndex(indcs::Tuple{Vararg{BlockIndex{1},N}}) where N = BlockIndex(block.(in
 ##
 
 @inline checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::Block{N}) where N = blockcheckbounds(Bool, A, I.n...)
+
 @inline function checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::BlockIndex{N}) where N
     bl = block(I)
     checkbounds(Bool, A, bl) || return false
-    B = A[bl]
-    checkbounds(Bool, B, blockindex(I)...)
+    # TODO: Replace with `eachblockaxes(A)[bl]` once that is defined.
+    binds = map(Base.axes1 ∘ getindex, axes(A), Tuple(bl))
+    Base.checkbounds_indices(Bool, binds, (blockindex(I),))
 end
-
-checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::AbstractVector{<:BlockIndex{N}}) where N =
-    all(checkbounds.(Bool, Ref(A), I))
+checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::AbstractArray{<:BlockIndex{N}}) where N =
+    all(i -> checkbounds(Bool, A, i), I)
 
 struct BlockIndexRange{N,R<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}},I<:Tuple{Vararg{Integer,N}},BI<:Integer} <: AbstractArray{BlockIndex{N,NTuple{N,BI},I},N}
     block::Block{N,BI}
@@ -243,6 +245,17 @@ length(iter::BlockIndexRange) = prod(size(iter))
 
 Block(bs::BlockIndexRange) = bs.block
 
+##
+# checkindex
+##
+
+function checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::BlockIndexRange{N}) where N
+    bl = block(I)
+    checkbounds(Bool, A, bl) || return false
+    # TODO: Replace with `eachblockaxes(A)[bl]` once that is defined.
+    binds = map(Base.axes1 ∘ getindex, axes(A), Tuple(bl))
+    Base.checkbounds_indices(Bool, binds, I.indices)
+end
 
 # #################
 # # support for pointers
@@ -432,6 +445,22 @@ _in(b, ::Tuple{}, ::Tuple{}, ::Tuple{}) = b
 # We sometimes need intersection of BlockRange to return a BlockRange
 intersect(a::BlockRange{1}, b::BlockRange{1}) = BlockRange((intersect(a.indices[1], b.indices[1]),))
 
+##
+# checkindex
+##
+
+# Used to ensure a `BlockBoundsError` is thrown instead of a `BoundsError`,
+# see https://github.com/JuliaArrays/BlockArrays.jl/issues/458
+checkbounds(A::AbstractArray{<:Any,N}, I::BlockRange{N}) where N = blockcheckbounds(A, I)
+checkbounds(A::AbstractArray, I1::BlockRange{1}, Irest::BlockRange{1}...) =
+    blockcheckbounds(A, I1, Irest...)
+
+# Convert Block inputs to integers.
+checkbounds(::Type{Bool}, A::AbstractArray{<:Any,N}, I::BlockRange{N}) where N =
+    blockcheckbounds(Bool, A, I.indices...)
+checkbounds(::Type{Bool}, A::AbstractArray, I1::AbstractVector{<:Block{1}}, Irest::AbstractVector{<:Block{1}}...) =
+    blockcheckbounds(Bool, A, map(I -> Int.(I), (I1, Irest...))...)
+
 # needed for scalar-like broadcasting
 
 BlockSlice{Block{1,BT},T,RT}(a::Base.OneTo) where {BT,T,RT<:AbstractUnitRange} =

test/test_blockarrays.jl

@@ -467,6 +467,60 @@ end
         BA_1 = BlockArray(undef_blocks, Vector{Float64}, [1,2,3])
         @test Base.IndexStyle(typeof(BA_1)) == IndexCartesian()
 
+        @test checkbounds(Bool, BA_1, Block(1))
+        @test isnothing(checkbounds(BA_1, Block(1)))
+        @test blockcheckbounds(Bool, BA_1, 1)
+        @test isnothing(blockcheckbounds(BA_1, 1))
+        @test checkbounds(Bool, BA_1, Block(2))
+        @test isnothing(checkbounds(BA_1, Block(2)))
+        @test blockcheckbounds(Bool, BA_1, 2)
+        @test isnothing(blockcheckbounds(BA_1, 2))
+        @test checkbounds(Bool, BA_1, Block(3))
+        @test isnothing(checkbounds(BA_1, Block(3)))
+        @test blockcheckbounds(Bool, BA_1, 3)
+        @test isnothing(blockcheckbounds(BA_1, 3))
+        @test !checkbounds(Bool, BA_1, Block(4))
+        @test_throws BlockBoundsError checkbounds(BA_1, Block(4))
+        @test !blockcheckbounds(Bool, BA_1, 4)
+        @test_throws BlockBoundsError blockcheckbounds(BA_1, 4)
+
+        @test checkbounds(Bool, BA_1, Block.(1:3))
+        @test isnothing(checkbounds(BA_1, Block.(1:3)))
+        @test blockcheckbounds(Bool, BA_1, 1:3)
+        @test isnothing(blockcheckbounds(BA_1, 1:3))
+        @test !checkbounds(Bool, BA_1, Block.(1:4))
+        @test_throws BlockBoundsError checkbounds(BA_1, Block.(1:4))
+        @test !blockcheckbounds(Bool, BA_1, 1:4)
+        @test_throws BlockBoundsError blockcheckbounds(BA_1, 1:4)
+
+        @test checkbounds(Bool, BA_1, [Block(1), Block(3)])
+        @test isnothing(checkbounds(BA_1, [Block(1), Block(3)]))
+        @test blockcheckbounds(Bool, BA_1, [1, 3])
+        @test isnothing(blockcheckbounds(BA_1, [1, 3]))
+        @test !checkbounds(Bool, BA_1, [Block(1), Block(4)])
+        @test_throws BlockBoundsError checkbounds(BA_1, [Block(1), Block(4)])
+        @test !blockcheckbounds(Bool, BA_1, [1, 4])
+        @test_throws BlockBoundsError blockcheckbounds(BA_1, [1, 4])
+
+        @test checkbounds(Bool, BA_1, Block(2)[2])
+        @test isnothing(checkbounds(BA_1, Block(2)[2]))
+        @test !checkbounds(Bool, BA_1, Block(2)[3])
+        @test_throws BoundsError checkbounds(BA_1, Block(2)[3])
+        @test !checkbounds(Bool, BA_1, Block(4)[2])
+        @test_throws BoundsError checkbounds(BA_1, Block(4)[2])
+
+        @test checkbounds(Bool, BA_1, Block(2)[1:2])
+        @test isnothing(checkbounds(BA_1, Block(2)[1:2]))
+        @test !checkbounds(Bool, BA_1, Block(2)[1:3])
+        @test_throws BoundsError checkbounds(BA_1, Block(2)[1:3])
+
+        @test checkbounds(Bool, BA_1, [Block(2)[2], Block(3)[3]])
+        @test isnothing(checkbounds(BA_1, [Block(2)[2], Block(3)[3]]))
+        @test !checkbounds(Bool, BA_1, [Block(2)[2], Block(2)[3]])
+        @test_throws BoundsError checkbounds(BA_1, [Block(2)[2], Block(2)[3]])
+        @test !checkbounds(Bool, BA_1, [Block(2)[2], Block(4)[2]])
+        @test_throws BoundsError checkbounds(BA_1, [Block(2)[2], Block(4)[2]])
+
         a_1 = rand(2)
         BA_1[Block(2)] = a_1
         @test BA_1[BlockIndex(2, 1)] == a_1[1]
@@ -491,6 +545,60 @@ end
         @test BA_2[Block(1,2)] == a_2
         BA_2[Block(1,2)] = a_2
 
+        @test checkbounds(Bool, BA_2, Block(1,1))
+        @test isnothing(checkbounds(BA_2, Block(1,1)))
+        @test checkbounds(Bool, BA_2, Block(1), Block(1))
+        @test isnothing(checkbounds(BA_2, Block(1), Block(1)))
+        @test !checkbounds(Bool, BA_2, Block(1,3))
+        @test_throws BlockBoundsError checkbounds(BA_2, Block(1,3))
+        @test !checkbounds(Bool, BA_2, Block(1), Block(3))
+        @test_throws BlockBoundsError checkbounds(BA_2, Block(1), Block(3))
+
+        @test checkbounds(Bool, BA_2, BlockRange(1:2,1:2))
+        @test blockcheckbounds(Bool, BA_2, BlockRange(1:2,1:2))
+        @test isnothing(checkbounds(BA_2, BlockRange(1:2,1:2)))
+        @test isnothing(blockcheckbounds(BA_2, BlockRange(1:2,1:2)))
+        @test checkbounds(Bool, BA_2, Block.(1:2), Block.(1:2))
+        @test checkbounds(Bool, BA_2, Block.([1,2]), Block.([1,2]))
+        @test isnothing(checkbounds(BA_2, Block.(1:2), Block.(1:2)))
+        @test isnothing(checkbounds(BA_2, Block.([1,2]), Block.([1,2])))
+        @test blockcheckbounds(Bool, BA_2, 1:2, 1:2)
+        @test blockcheckbounds(Bool, BA_2, [1,2], [1,2])
+        @test isnothing(blockcheckbounds(BA_2, 1:2, 1:2))
+        @test isnothing(blockcheckbounds(BA_2, Block.(1:2), Block.(1:2)))
+        @test isnothing(blockcheckbounds(BA_2, [1,2], [1,2]))
+        @test isnothing(blockcheckbounds(BA_2, Block.([1,2]), Block.([1,2])))
+
+        @test !checkbounds(Bool, BA_2, BlockRange(1:2,1:3))
+        @test !blockcheckbounds(Bool, BA_2, BlockRange(1:2,1:3))
+        @test_throws BlockBoundsError checkbounds(BA_2, BlockRange(1:2,1:3))
+        @test_throws BlockBoundsError blockcheckbounds(BA_2, BlockRange(1:2,1:3))
+        @test !checkbounds(Bool, BA_2, Block.(1:2), Block.(1:3))
+        @test !checkbounds(Bool, BA_2, Block.([1,2]), Block.([1,3]))
+        @test_throws BlockBoundsError checkbounds(BA_2, Block.(1:2), Block.(1:3))
+        @test_throws BlockBoundsError checkbounds(BA_2, Block.([1,2]), Block.([1,3]))
+        @test !blockcheckbounds(Bool, BA_2, 1:2, 1:3)
+        @test !blockcheckbounds(Bool, BA_2, [1,2], [1,3])
+        @test_throws BlockBoundsError blockcheckbounds(BA_2, 1:2, 1:3)
+        @test_throws BlockBoundsError blockcheckbounds(BA_2, Block.(1:2), Block.(1:3))
+        @test_throws BlockBoundsError blockcheckbounds(BA_2, [1,2], [1,3])
+        @test_throws BlockBoundsError blockcheckbounds(BA_2, Block.([1,2]), Block.([1,3]))
+
+        @test checkbounds(Bool, BA_2, Block(1,1)[1,1])
+        @test isnothing(checkbounds(BA_2, Block(1,1)[1,1]))
+        @test checkbounds(Bool, BA_2, Block(1)[1], Block(1)[1])
+        @test checkbounds(Bool, BA_2, Block(1,2)[1,2])
+        @test !checkbounds(Bool, BA_2, Block(1,2)[2,2])
+        @test checkbounds(Bool, BA_2, Block(2,2)[1:1,1:2])
+        @test checkbounds(Bool, BA_2, Matrix(Block(2,2)[1:1,1:2]))
+        @test isnothing(checkbounds(BA_2, Matrix(Block(2,2)[1:1,1:2])))
+        @test checkbounds(Bool, BA_2, Block(1,2)[1:1,1:2])
+        @test checkbounds(Bool, BA_2, Block(1)[1:1], Block(2)[1:2])
+        @test !checkbounds(Bool, BA_2, Block(2,2)[1:3,1:2])
+        @test !checkbounds(Bool, BA_2, Block(2)[1:3], Block(2)[1:2])
+        @test checkbounds(Bool, BA_2, [Block(2)[1], Block(2)[2]], Block(2)[1:2])
+        @test !checkbounds(Bool, BA_2, [Block(2)[1], Block(2)[3]], Block(2)[1:2])
+
         @test BA_2[1,5] == a_2[2]
         @test_throws DimensionMismatch BA_2[Block(1,2)] = rand(1,5)
 

test/test_blockindices.jl

@@ -380,6 +380,14 @@ end
         @test !checkbounds(Bool, b, Block(3)[4])
         @test !checkbounds(Bool, b, Block(0)[1])
         @test !checkbounds(Bool, b, Block(1)[0])
+        @test checkbounds(Bool, b, Block(1))
+        @test checkbounds(Bool, b, Block(2))
+        @test checkbounds(Bool, b, Block(3))
+        @test !checkbounds(Bool, b, Block(4))
+        @test checkbounds(Bool, b, Block.(1:3))
+        @test !checkbounds(Bool, b, Block.(1:4))
+        @test_throws BlockBoundsError checkbounds(b, Block(4))
+        @test_throws BlockBoundsError checkbounds(b, Block.(1:4))
         # treat b as the axis
         @test checkindex(Bool, b, Block(1)[1])
         @test checkindex(Bool, b, Block(1)[1:1])