Add findblockindex support for AbstractArray

src/abstractblockarray.jl

@@ -157,6 +157,16 @@
 """
 @inline blockcheckindex(::Type{Bool}, inds::BlockRange{1}, i::Integer) = Block(i) in inds
 
+"""
+    to_blockindex(A::AbstractArray{T,N}, I::Tuple{Vararg{Integer,N}})
+
+Returns the [`BlockIndex`](@ref) of `A[I...]`.
+"""
+function to_blockindex(A::AbstractArray{T,N}, I::Tuple{Vararg{Integer,N}}) where {T,N}
+    blockinds = _to_blockindex.(axes(A), I)
+    BlockIndex(map(first,blockinds), map(last,blockinds))
+end
+
 @propagate_inbounds setindex!(block_arr::AbstractBlockArray{T,N}, v, block::Block{N}) where {T,N} =
     setindex!(block_arr, v, Block.(block.n)...)
 @propagate_inbounds function setindex!(block_arr::AbstractBlockArray{T,N}, v, block::Vararg{Block{1}, N}) where {T,N}

src/blockaxis.jl

@@ -184,6 +184,18 @@
 @inline blockedrange(blocks::Union{Tuple,AbstractVector}) = BlockedOneTo(_blocklengths2blocklasts(blocks))
 @inline blockedrange(f::Integer, blocks::Union{Tuple,AbstractVector}) = _BlockedUnitRange(f, f-oneunit(f) .+ _blocklengths2blocklasts(blocks))
 
+to_blockindex(b::BlockedUnitRange, k::Integer) = BlockIndex(_to_blockindex(b,k)...)
+function _to_blockindex(b::BlockedUnitRange, k::Integer)
+    if isone(b.first)
+        bl = blocklasts(b)
+        prevblocklast = k == firstindex(bl) ? first(b)-1 : bl[k-1]
+        local_index = k - prevblocklast
+        (k, local_index)
+    else
+        _to_blockindex(axes(b,1), k)
+    end
+end
+
 _diff(a::AbstractVector) = diff(a)
 _diff(a::Tuple) = diff(collect(a))
 @inline _blocklengths(a, bl, dbl) = isempty(bl) ? [dbl;] : [first(bl)-first(a)+oneunit(eltype(a)); dbl]