diff --git a/src/abstractblockarray.jl b/src/abstractblockarray.jl index e16bf42e..ffce680c 100644 --- a/src/abstractblockarray.jl +++ b/src/abstractblockarray.jl @@ -157,6 +157,16 @@ false """ @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} diff --git a/src/blockaxis.jl b/src/blockaxis.jl index 75ab26de..a6d3dee0 100644 --- a/src/blockaxis.jl +++ b/src/blockaxis.jl @@ -184,6 +184,18 @@ julia> blockedrange(2, (1,2)) @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]