From 8757cca48bedc7e76784002f5b0f11ec3718f01a Mon Sep 17 00:00:00 2001
From: Tianyi Pu <44583944+putianyi889@users.noreply.github.com>
Date: Thu, 21 Mar 2024 11:40:24 +0000
Subject: [PATCH 1/2] Update abstractblockarray.jl

---
 src/abstractblockarray.jl | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/abstractblockarray.jl b/src/abstractblockarray.jl
index e16bf42e..76fe2b9c 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
 
+"""
+    findblockindex(A::AbstractArray{T,N}, I::Tuple{Vararg{Integer,N}})
+
+Returns the [`BlockIndex`](@ref) of `A[I...]`.
+"""
+function findblockindex(A::AbstractArray{T,N}, I::Tuple{Vararg{Integer,N}}) where {T,N}
+    blockinds = findblockindex.(axes(A), I)
+    BlockIndex(Base.IteratorsMD.flatten(map(x->x.I, blockinds)), Base.IteratorsMD.flatten(map(x->x.α, 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}

From bd8b591b31e29c7dd32334086d7c69df0f3cb870 Mon Sep 17 00:00:00 2001
From: Tianyi Pu <912396513@qq.com>
Date: Sun, 31 Mar 2024 23:16:14 +0100
Subject: [PATCH 2/2] change name to `to_blockindex`

---
 src/abstractblockarray.jl |  8 ++++----
 src/blockaxis.jl          | 14 +++++++++++++-
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/abstractblockarray.jl b/src/abstractblockarray.jl
index 76fe2b9c..ffce680c 100644
--- a/src/abstractblockarray.jl
+++ b/src/abstractblockarray.jl
@@ -158,13 +158,13 @@ false
 @inline blockcheckindex(::Type{Bool}, inds::BlockRange{1}, i::Integer) = Block(i) in inds
 
 """
-    findblockindex(A::AbstractArray{T,N}, I::Tuple{Vararg{Integer,N}})
+    to_blockindex(A::AbstractArray{T,N}, I::Tuple{Vararg{Integer,N}})
 
 Returns the [`BlockIndex`](@ref) of `A[I...]`.
 """
-function findblockindex(A::AbstractArray{T,N}, I::Tuple{Vararg{Integer,N}}) where {T,N}
-    blockinds = findblockindex.(axes(A), I)
-    BlockIndex(Base.IteratorsMD.flatten(map(x->x.I, blockinds)), Base.IteratorsMD.flatten(map(x->x.α, blockinds)))
+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} =
diff --git a/src/blockaxis.jl b/src/blockaxis.jl
index de144392..3961915e 100644
--- a/src/blockaxis.jl
+++ b/src/blockaxis.jl
@@ -15,7 +15,7 @@ function findblockindex(b::AbstractVector, k::Integer)
     @boundscheck k in b || throw(BoundsError())
     bl = blocklasts(b)
     blockidx = _searchsortedfirst(bl, k)
-    @assert blockindex != lastindex(bl) + 1 # guaranteed by the @boundscheck above
+    @assert blockindex != lastindex(bl) + 1  # guaranteed by the @boundscheck above
     prevblocklast = blockidx == firstindex(bl) ? first(b)-1 : bl[blockidx-1]
     local_index = k - prevblocklast
     return BlockIndex(blockidx, local_index)
@@ -72,6 +72,18 @@ _blocklengths2blocklasts(blocks) = cumsum(blocks) # extra level to allow changin
 end
 @inline blocklasts(a::BlockedUnitRange) = a.lasts
 
+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::BlockedUnitRange) = isempty(blocklasts(a)) ? [_diff(blocklasts(a));] : [first(blocklasts(a))-first(a)+1; _diff(blocklasts(a))]