Skip to content

Commit 5a8301c

Browse files
lkdvoslkdvos
committed
Rework indexing as a whole
1 parent 9e9e1f0 commit 5a8301c

File tree

5 files changed

+332
-130
lines changed

5 files changed

+332
-130
lines changed

‎src/SparseArraysBase.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ export SparseArrayDOK,
1313
storedpairs,
1414
storedvalues
1515

16+
include("indexing.jl")
1617
include("abstractsparsearrayinterface.jl")
1718
include("sparsearrayinterface.jl")
1819
include("wrappers.jl")

‎src/abstractsparsearray.jl‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@ using DerivableInterfaces: @array_aliases
55
@array_aliases AbstractSparseArray
66

77
using DerivableInterfaces: DerivableInterfaces
8-
function DerivableInterfaces.interface(::Type{<:AbstractSparseArray})
9-
return SparseArrayInterface()
8+
function DerivableInterfaces.interface(::Type{A}) where {A<:AnyAbstractSparseArray}
9+
return SparseArrayInterface{ndims(A)}()
1010
end
1111

1212
using DerivableInterfaces: @derive

‎src/abstractsparsearrayinterface.jl‎

Lines changed: 99 additions & 126 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,6 @@ using DerivableInterfaces: DerivableInterfaces, @derive, @interface, AbstractArr
44
# since it is considered part of the sparse array interface.
55
using ArrayLayouts: zero!
66

7-
function eachstoredindex end
8-
function getstoredindex end
9-
function getunstoredindex end
10-
function isstored end
11-
function setstoredindex! end
12-
function setunstoredindex! end
13-
function storedlength end
14-
function storedpairs end
15-
function storedvalues end
16-
177
# Replace the function for accessing
188
# unstored values.
199
function set_getunstoredindex end
@@ -34,79 +24,23 @@ end
3424

3525
# Minimal interface for `SparseArrayInterface`.
3626
# Fallbacks for dense/non-sparse arrays.
37-
@interface ::AbstractArrayInterface isstored(a::AbstractArray, I::Int...) = true
38-
@interface ::AbstractArrayInterface eachstoredindex(a::AbstractArray) = eachindex(a)
39-
@interface ::AbstractArrayInterface getstoredindex(a::AbstractArray, I::Int...) =
40-
getindex(a, I...)
41-
@interface ::AbstractArrayInterface function setstoredindex!(
42-
a::AbstractArray, value, I::Int...
43-
)
44-
setindex!(a, value, I...)
45-
return a
46-
end
47-
# TODO: Should this error by default if the value at the index
48-
# is stored? It could be disabled with something analogous
49-
# to `checkbounds`, like `checkstored`/`checkunstored`.
50-
@interface ::AbstractArrayInterface function setunstoredindex!(
51-
a::AbstractArray, value, I::Int...
52-
)
53-
# TODO: Make this a `MethodError`?
54-
return error("Not implemented.")
55-
end
56-
57-
# TODO: Use `Base.to_indices`?
58-
isstored(a::AbstractArray, I::CartesianIndex) = isstored(a, Tuple(I)...)
59-
# TODO: Use `Base.to_indices`?
60-
getstoredindex(a::AbstractArray, I::CartesianIndex) = getstoredindex(a, Tuple(I)...)
61-
# TODO: Use `Base.to_indices`?
62-
getunstoredindex(a::AbstractArray, I::CartesianIndex) = getunstoredindex(a, Tuple(I)...)
63-
# TODO: Use `Base.to_indices`?
64-
function setstoredindex!(a::AbstractArray, value, I::CartesianIndex)
65-
return setstoredindex!(a, value, Tuple(I)...)
66-
end
67-
# TODO: Use `Base.to_indices`?
68-
function setunstoredindex!(a::AbstractArray, value, I::CartesianIndex)
69-
return setunstoredindex!(a, value, Tuple(I)...)
70-
end
71-
72-
# Interface defaults.
73-
# TODO: Have a fallback that handles element types
74-
# that don't define `zero(::Type)`.
75-
@interface ::AbstractArrayInterface getunstoredindex(a::AbstractArray, I::Int...) =
76-
zero(eltype(a))
77-
78-
# DerivableInterfacesd interface.
79-
@interface ::AbstractArrayInterface storedlength(a::AbstractArray) = length(storedvalues(a))
80-
@interface ::AbstractArrayInterface storedpairs(a::AbstractArray) =
81-
map(I -> I => getstoredindex(a, I), eachstoredindex(a))
82-
83-
@interface ::AbstractArrayInterface function eachstoredindex(as::AbstractArray...)
84-
return eachindex(as...)
85-
end
86-
87-
@interface ::AbstractArrayInterface storedvalues(a::AbstractArray) = a
88-
89-
# Automatically derive the interface for all `AbstractArray` subtypes.
90-
# TODO: Define `SparseArrayInterfaceOps` derivable trait and rewrite this
91-
# as `@derive AbstractArray SparseArrayInterfaceOps`.
92-
@derive (T=AbstractArray,) begin
93-
SparseArraysBase.eachstoredindex(::T)
94-
SparseArraysBase.eachstoredindex(::T...)
95-
SparseArraysBase.getstoredindex(::T, ::Int...)
96-
SparseArraysBase.getunstoredindex(::T, ::Int...)
97-
SparseArraysBase.isstored(::T, ::Int...)
98-
SparseArraysBase.setstoredindex!(::T, ::Any, ::Int...)
99-
SparseArraysBase.setunstoredindex!(::T, ::Any, ::Int...)
100-
SparseArraysBase.storedlength(::T)
101-
SparseArraysBase.storedpairs(::T)
102-
SparseArraysBase.storedvalues(::T)
103-
end
104-
105-
# TODO: Add `ndims` type parameter, like `Base.Broadcast.AbstractArrayStyle`.
106-
# TODO: This isn't used to define interface functions right now.
107-
# Currently, `@interface` expects an instance, probably it should take a
108-
# type instead so fallback functions can use abstract types.
109-
abstract type AbstractSparseArrayInterface <: AbstractArrayInterface end
27+
28+
"""
29+
AbstractSparseArrayInterface{N} <: AbstractArrayInterface{N}
30+
31+
Abstract supertype for any interface associated with sparse array types.
32+
"""
33+
abstract type AbstractSparseArrayInterface{N} <: AbstractArrayInterface{N} end
34+
35+
"""
36+
SparseArrayInterface{N} <: AbstractSparseArrayInterface{N}
37+
38+
Interface for array operations that are centered around sparse storage types, typically assuming
39+
fast `O(1)` random access/insertion, but slower sequential access.
40+
"""
41+
struct SparseArrayInterface <: AbstractSparseArrayInterface end
42+
43+
# by default, this interface is stronger than other interfaces (is this fair?)
11044

11145
function DerivableInterfaces.combine_interface_rule(
11246
interface1::AbstractSparseArrayInterface, interface2::AbstractSparseArrayInterface
@@ -129,10 +63,88 @@ function DerivableInterfaces.combine_interface_rule(
12963
return interface2
13064
end
13165

132-
# using ArrayLayouts getindex: this is a bit cumbersome because there already is a way to make that work focussed on types
133-
# but here we want to focus on interfaces.
134-
# eg: ArrayLayouts.@layoutgetindex ArrayType
135-
# TODO: decide if we need the interface approach,
66+
# getindex/setindex!
67+
# ------------------
68+
# canonical errors are moved to `isstored`, `getstoredindex` and `getunstoredindex`
69+
# so no errors at this level by defining both IndexLinear and IndexCartesian
70+
@interface ::AbstractSparseArrayInterface function Base.getindex(
71+
A::AbstractArray{<:Any,N}, I::Vararg{Int,N}
72+
) where {N}
73+
@_propagate_inbounds_meta
74+
@boundscheck checkbounds(A, I...) # generally isstored requires bounds checking
75+
return @inbounds isstored(A, I...) ? getstoredindex(A, I...) : getunstoredindex(A, I...)
76+
end
77+
@interface ::AbstractSparseArrayInterface function Base.getindex(A::AbstractArray, I::Int)
78+
@_propagate_inbounds_meta
79+
@boundscheck checkbounds(A, I)
80+
return @inbounds isstored(A, I) ? getstoredindex(A, I) : getunstoredindex(A, I)
81+
end
82+
# disambiguate vectors
83+
@interface ::AbstractSparseArrayInterface function Base.getindex(A::AbstractVector, I::Int)
84+
@_propagate_inbounds_meta
85+
@boundscheck checkbounds(A, I)
86+
return @inbounds isstored(A, I) ? getstoredindex(A, I) : getunstoredindex(A, I)
87+
end
88+
89+
@interface ::AbstractSparseArrayInterface function Base.setindex!(
90+
A::AbstractArray{<:Any,N}, v, I::Vararg{Int,N}
91+
) where {N}
92+
@_propagate_inbounds_meta
93+
@boundscheck checkbounds(A, I...)
94+
return @inbounds if isstored(A, I...)
95+
setstoredindex!(A, v, I...)
96+
else
97+
setunstoredindex!(A, v, I...)
98+
end
99+
end
100+
@interface ::AbstractSparseArrayInterface function Base.setindex!(A::AbstractArray, I::Int)
101+
@_propagate_inbounds_meta
102+
@boundscheck checkbounds(A, I)
103+
return @inbounds if isstored(A, I)
104+
setstoredindex!(A, v, I)
105+
else
106+
setunstoredindex!(A, v, I)
107+
end
108+
end
109+
# disambiguate vectors
110+
@interface ::AbstractSparseArrayInterface function Base.setindex!(A::AbstractVector, I::Int)
111+
@_propagate_inbounds_meta
112+
@boundscheck checkbounds(A, I)
113+
return @inbounds if isstored(A, I)
114+
setstoredindex!(A, v, I)
115+
else
116+
setunstoredindex!(A, v, I)
117+
end
118+
end
119+
120+
# Indices
121+
# -------
122+
# required:
123+
@interface ::AbstractSparseArrayInterface eachstoredindex(::IndexStyle, A::AbstractArray) =
124+
throw(MethodError(eachstoredindex, (style, A)))
125+
@interface ::AbstractSparseArrayInterface storedvalues(A::AbstractArray) =
126+
throw(MethodError(storedvalues, A))
127+
128+
# derived but may be specialized:
129+
@interface ::AbstractSparseArrayInterface function eachstoredindex(
130+
style::IndexStyle, A::AbstractArray, B::AbstractArray...
131+
)
132+
return union(map(Base.Fix1(eachstoredindex, style), (A, B...))...)
133+
end
134+
135+
@interface ::AbstractSparseArrayInterface storedlength(A::AbstractArray) =
136+
length(storedvalues(A))
137+
@interface ::AbstractSparseArrayInterface storedpairs(A::AbstractArray) =
138+
zip(eachstoredindex(A), storedvalues(A))
139+
140+
#=
141+
All sparse array interfaces are mapped through layout_getindex. (is this too opinionated?)
142+
143+
using ArrayLayouts getindex: this is a bit cumbersome because there already is a way to make that work focussed on types
144+
but here we want to focus on interfaces.
145+
eg: ArrayLayouts.@layoutgetindex ArrayType
146+
TODO: decide if we need the interface approach at all here
147+
=#
136148
for (Tr, Tc) in Iterators.product(
137149
Iterators.repeated((:Colon, :AbstractUnitRange, :AbstractVector, :Integer), 2)...
138150
)
@@ -172,45 +184,6 @@ end
172184

173185
@interface ::AbstractSparseArrayInterface storedvalues(a::AbstractArray) = StoredValues(a)
174186

175-
@interface ::AbstractSparseArrayInterface function eachstoredindex(
176-
a1::AbstractArray, a2::AbstractArray, a_rest::AbstractArray...
177-
)
178-
# TODO: Make this more customizable, say with a function
179-
# `combine/promote_storedindices(a1, a2)`.
180-
return union(eachstoredindex.((a1, a2, a_rest...))...)
181-
end
182-
183-
@interface ::AbstractSparseArrayInterface function eachstoredindex(a::AbstractArray)
184-
# TODO: Use `MethodError`?
185-
return error("Not implemented.")
186-
end
187-
188-
# We restrict to `I::Vararg{Int,N}` to allow more general functions to handle trailing
189-
# indices and linear indices.
190-
@interface ::AbstractSparseArrayInterface function Base.getindex(
191-
a::AbstractArray{<:Any,N}, I::Vararg{Int,N}
192-
) where {N}
193-
!isstored(a, I...) && return getunstoredindex(a, I...)
194-
return getstoredindex(a, I...)
195-
end
196-
197-
# We restrict to `I::Vararg{Int,N}` to allow more general functions to handle trailing
198-
# indices and linear indices.
199-
@interface ::AbstractSparseArrayInterface function Base.setindex!(
200-
a::AbstractArray{<:Any,N}, value, I::Vararg{Int,N}
201-
) where {N}
202-
if !isstored(a, I...)
203-
# Don't set the value if it is zero, but only check
204-
# if it is zero if the elements are numbers since otherwise
205-
# it may be nontrivial to check.
206-
eltype(a) <: Number && iszero(value) && return a
207-
setunstoredindex!(a, value, I...)
208-
return a
209-
end
210-
setstoredindex!(a, value, I...)
211-
return a
212-
end
213-
214187
# TODO: This may need to be defined in `sparsearraydok.jl`, after `SparseArrayDOK`
215188
# is defined. And/or define `default_type(::SparseArrayStyle, T::Type) = SparseArrayDOK{T}`.
216189
@interface ::AbstractSparseArrayInterface function Base.similar(

0 commit comments

Comments
 (0)