Skip to content

Commit bd54ec3

Browse files
committed
Fix some tests
1 parent a3d0b91 commit bd54ec3

File tree

1 file changed

+29
-16
lines changed

1 file changed

+29
-16
lines changed

src/BlockArraysExtensions/BlockArraysExtensions.jl

Lines changed: 29 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -21,31 +21,39 @@ using BlockArrays:
2121
findblockindex
2222
using Dictionaries: Dictionary, Indices
2323
using GradedUnitRanges: blockedunitrange_getindices, to_blockindices
24-
using SparseArraysBase: SparseArraysBase, storedlength, eachstoredindex
24+
using SparseArraysBase:
25+
SparseArraysBase,
26+
eachstoredindex,
27+
getunstoredindex,
28+
isstored,
29+
setunstoredindex!,
30+
storedlength
2531

2632
# A return type for `blocks(array)` when `array` isn't blocked.
2733
# Represents a vector with just that single block.
2834
struct SingleBlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N}
2935
array::Array
3036
end
37+
Base.parent(a::SingleBlockView) = a.array
3138
blocks_maybe_single(a) = blocks(a)
3239
blocks_maybe_single(a::Array) = SingleBlockView(a)
3340
function Base.getindex(a::SingleBlockView{<:Any,N}, index::Vararg{Int,N}) where {N}
3441
@assert all(isone, index)
35-
return a.array
42+
return parent(a)
3643
end
3744

3845
# A wrapper around a potentially blocked array that is not blocked.
3946
struct NonBlockedArray{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N}
4047
array::Array
4148
end
42-
Base.size(a::NonBlockedArray) = size(a.array)
43-
Base.getindex(a::NonBlockedArray{<:Any,N}, I::Vararg{Integer,N}) where {N} = a.array[I...]
49+
Base.parent(a::NonBlockedArray) = a.array
50+
Base.size(a::NonBlockedArray) = size(parent(a))
51+
Base.getindex(a::NonBlockedArray{<:Any,N}, I::Vararg{Integer,N}) where {N} = parent(a)[I...]
4452
# Views of `NonBlockedArray`/`NonBlockedVector` are eager.
4553
# This fixes an issue in Julia 1.11 where reindexing defaults to using views.
4654
# TODO: Maybe reconsider this design, and allows views to work in slicing.
4755
Base.view(a::NonBlockedArray, I...) = a[I...]
48-
BlockArrays.blocks(a::NonBlockedArray) = SingleBlockView(a.array)
56+
BlockArrays.blocks(a::NonBlockedArray) = SingleBlockView(parent(a))
4957
const NonBlockedVector{T,Array} = NonBlockedArray{T,1,Array}
5058
NonBlockedVector(array::AbstractVector) = NonBlockedArray(array)
5159

@@ -100,10 +108,10 @@ Base.view(S::BlockIndices, i) = S[i]
100108
function Base.getindex(
101109
a::NonBlockedVector{<:Integer,<:BlockIndices}, I::UnitRange{<:Integer}
102110
)
103-
ax = only(axes(a.array.indices))
111+
ax = only(axes(parent(a).indices))
104112
brs = to_blockindices(ax, I)
105113
inds = blockedunitrange_getindices(ax, I)
106-
return NonBlockedVector(a.array[BlockSlice(brs, inds)])
114+
return NonBlockedVector(parent(a)[BlockSlice(brs, inds)])
107115
end
108116

109117
function Base.getindex(S::BlockIndices, i::BlockSlice{<:BlockRange{1}})
@@ -511,25 +519,30 @@ struct BlockView{T,N,Array<:AbstractArray{T,N}} <: AbstractArray{T,N}
511519
array::Array
512520
block::Tuple{Vararg{Block{1,Int},N}}
513521
end
522+
Base.parent(a::BlockView) = a.array
514523
function Base.axes(a::BlockView)
515524
# TODO: Try to avoid conversion to `Base.OneTo{Int}`, or just convert
516525
# the element type to `Int` with `Int.(...)`.
517-
# When the axes of `a.array` are `GradedOneTo`, the block is `LabelledUnitRange`,
526+
# When the axes of `parent(a)` are `GradedOneTo`, the block is `LabelledUnitRange`,
518527
# which has element type `LabelledInteger`. That causes conversion problems
519528
# in some generic Base Julia code, for example when printing `BlockView`.
520529
return ntuple(ndims(a)) do dim
521-
return Base.OneTo{Int}(only(axes(axes(a.array, dim)[a.block[dim]])))
530+
return Base.OneTo{Int}(only(axes(axes(parent(a), dim)[a.block[dim]])))
522531
end
523532
end
524533
function Base.size(a::BlockView)
525534
return length.(axes(a))
526535
end
527536
function Base.getindex(a::BlockView{<:Any,N}, index::Vararg{Int,N}) where {N}
528-
return blocks(a.array)[Int.(a.block)...][index...]
537+
return blocks(parent(a))[Int.(a.block)...][index...]
529538
end
530539
function Base.setindex!(a::BlockView{<:Any,N}, value, index::Vararg{Int,N}) where {N}
531-
blocks(a.array)[Int.(a.block)...] = blocks(a.array)[Int.(a.block)...]
532-
blocks(a.array)[Int.(a.block)...][index...] = value
540+
I = Int.(a.block)
541+
if !isstored(blocks(parent(a)), I...)
542+
unstored_value = getunstoredindex(blocks(parent(a)), I...)
543+
setunstoredindex!(blocks(parent(a)), unstored_value, I...)
544+
end
545+
blocks(parent(a))[I...][index...] = value
533546
return a
534547
end
535548

@@ -538,15 +551,15 @@ function SparseArraysBase.storedlength(a::BlockView)
538551
# a Bool in `BlockView`.
539552
I = CartesianIndex(Int.(a.block))
540553
# TODO: Use `block_eachstoredindex`.
541-
if I eachstoredindex(blocks(a.array))
542-
return storedlength(blocks(a.array)[I])
554+
if I eachstoredindex(blocks(parent(a)))
555+
return storedlength(blocks(parent(a))[I])
543556
end
544557
return 0
545558
end
546559

547560
## # Allow more fine-grained control:
548561
## function ArrayLayouts.sub_materialize(layout, a::BlockView, ax)
549-
## return blocks(a.array)[Int.(a.block)...]
562+
## return blocks(parent(a))[Int.(a.block)...]
550563
## end
551564
## function ArrayLayouts.sub_materialize(layout, a::BlockView)
552565
## return sub_materialize(layout, a, axes(a))
@@ -555,7 +568,7 @@ end
555568
## return sub_materialize(MemoryLayout(a), a)
556569
## end
557570
function ArrayLayouts.sub_materialize(a::BlockView)
558-
return blocks(a.array)[Int.(a.block)...]
571+
return blocks(parent(a))[Int.(a.block)...]
559572
end
560573

561574
function view!(a::AbstractArray{<:Any,N}, index::Block{N}) where {N}

0 commit comments

Comments
 (0)