|
1 | 1 | using BlockArrays: Block
|
2 | 2 | using DerivableInterfaces: zero!
|
3 | 3 |
|
4 |
| -struct GetUnstoredBlock{Axes} |
5 |
| - axes::Axes |
| 4 | +struct ZeroBlocks{ |
| 5 | + N,A<:AbstractArray{<:Any,N},ParentAxes<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}} |
| 6 | +} <: AbstractArray{A,N} |
| 7 | + parentaxes::ParentAxes |
| 8 | +end |
| 9 | +function ZeroBlocks{N,A}( |
| 10 | + ax::Ax |
| 11 | +) where {N,A<:AbstractArray{<:Any,N},Ax<:Tuple{Vararg{AbstractUnitRange{<:Integer},N}}} |
| 12 | + return ZeroBlocks{N,A,Ax}(ax) |
| 13 | +end |
| 14 | +Base.size(a::ZeroBlocks) = map(blocklength, a.parentaxes) |
| 15 | + |
| 16 | +function Base.AbstractArray{A}(a::ZeroBlocks{N}) where {N,A} |
| 17 | + return ZeroBlocks{N,A}(a.parentaxes) |
6 | 18 | end
|
7 | 19 |
|
8 |
| -@inline function (f::GetUnstoredBlock)( |
9 |
| - ::Type{<:AbstractArray{A,N}}, I::Vararg{Int,N} |
10 |
| -) where {A,N} |
| 20 | +@inline function Base.getindex(a::ZeroBlocks{N,A}, I::Vararg{Int,N}) where {N,A} |
| 21 | + # TODO: Use `BlockArrays.eachblockaxes`. |
11 | 22 | ax = ntuple(N) do d
|
12 |
| - return only(axes(f.axes[d][Block(I[d])])) |
| 23 | + return only(axes(a.parentaxes[d][Block(I[d])])) |
13 | 24 | end
|
14 | 25 | !isconcretetype(A) && return zero!(similar(Array{eltype(A),N}, ax))
|
15 | 26 | return zero!(similar(A, ax))
|
16 | 27 | end
|
17 |
| -@inline function (f::GetUnstoredBlock)( |
18 |
| - a::AbstractArray{<:Any,N}, I::Vararg{Int,N} |
19 |
| -) where {N} |
20 |
| - return f(typeof(a), I...) |
21 |
| -end |
22 | 28 | # TODO: Use `Base.to_indices`.
|
23 |
| -@inline function (f::GetUnstoredBlock)( |
24 |
| - a::AbstractArray{<:Any,N}, I::CartesianIndex{N} |
25 |
| -) where {N} |
26 |
| - return f(a, Tuple(I)...) |
| 29 | +@inline function Base.getindex(a::ZeroBlocks{N,A}, I::CartesianIndex{N}) where {N,A} |
| 30 | + return a[Tuple(I)...] |
27 | 31 | end
|
28 | 32 |
|
29 | 33 | # TODO: this is a hack and is also type-unstable
|
30 | 34 | using LinearAlgebra: Diagonal
|
31 | 35 | using TypeParameterAccessors: similartype
|
32 |
| -function (f::GetUnstoredBlock)( |
33 |
| - ::Type{<:AbstractMatrix{<:Diagonal{<:Any,V}}}, I::Vararg{Int,2} |
34 |
| -) where {V} |
| 36 | +function Base.getindex(a::ZeroBlocks{2,A}, I::Vararg{Int,2}) where {V,A<:Diagonal{<:Any,V}} |
35 | 37 | ax = ntuple(2) do d
|
36 |
| - return only(axes(f.axes[d][Block(I[d])])) |
| 38 | + return only(axes(a.parentaxes[d][Block(I[d])])) |
37 | 39 | end
|
38 | 40 | if allequal(I)
|
39 | 41 | return Diagonal(zero!(similar(V, first(ax))))
|
|
0 commit comments