Skip to content

Commit 461a3d0

Browse files
lkdvoslkdvos
committed
Clean up indexing
1 parent 2251dd4 commit 461a3d0

File tree

3 files changed

+120
-192
lines changed

3 files changed

+120
-192
lines changed

‎src/abstractsparsearrayinterface.jl‎

Lines changed: 25 additions & 166 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,6 @@ function densearray(a::AbstractArray)
2121
return Array(a)
2222
end
2323

24-
# Minimal interface for `SparseArrayInterface`.
25-
# Fallbacks for dense/non-sparse arrays.
26-
2724
"""
2825
AbstractSparseArrayInterface <: AbstractArrayInterface
2926
@@ -62,100 +59,21 @@ function DerivableInterfaces.combine_interface_rule(
6259
return interface2
6360
end
6461

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

16179
to_vec(x) = vec(collect(x))
@@ -171,18 +89,18 @@ to_vec(x::AbstractArray) = vec(x)
17189
# Most sparse arrays should overload `storedvalues` directly
17290
# and avoid this wrapper since it adds extra indirection to
17391
# access stored values.
174-
struct StoredValues{T,A<:AbstractArray{T},I} <: AbstractVector{T}
175-
array::A
176-
storedindices::I
177-
end
178-
StoredValues(a::AbstractArray) = StoredValues(a, to_vec(eachstoredindex(a)))
179-
Base.size(a::StoredValues) = size(a.storedindices)
180-
Base.getindex(a::StoredValues, I::Int) = getstoredindex(a.array, a.storedindices[I])
181-
function Base.setindex!(a::StoredValues, value, I::Int)
182-
return setstoredindex!(a.array, value, a.storedindices[I])
183-
end
184-
185-
@interface ::AbstractSparseArrayInterface storedvalues(a::AbstractArray) = StoredValues(a)
92+
# struct StoredValues{T,A<:AbstractArray{T},I} <: AbstractVector{T}
93+
# array::A
94+
# storedindices::I
95+
# end
96+
# StoredValues(a::AbstractArray) = StoredValues(a, to_vec(eachstoredindex(a)))
97+
# Base.size(a::StoredValues) = size(a.storedindices)
98+
# Base.getindex(a::StoredValues, I::Int) = getstoredindex(a.array, a.storedindices[I])
99+
# function Base.setindex!(a::StoredValues, value, I::Int)
100+
# return setstoredindex!(a.array, value, a.storedindices[I])
101+
# end
102+
#
103+
# @interface ::AbstractSparseArrayInterface storedvalues(a::AbstractArray) = StoredValues(a)
186104

187105
# TODO: This may need to be defined in `sparsearraydok.jl`, after `SparseArrayDOK`
188106
# is defined. And/or define `default_type(::SparseArrayStyle, T::Type) = SparseArrayDOK{T}`.
@@ -193,38 +111,6 @@ end
193111
return SparseArrayDOK{T}(size...)
194112
end
195113

196-
# map over a specified subset of indices of the inputs.
197-
function map_indices! end
198-
199-
@interface interface::AbstractArrayInterface function map_indices!(
200-
indices, f, a_dest::AbstractArray, as::AbstractArray...
201-
)
202-
for I in indices
203-
a_dest[I] = f(map(a -> a[I], as)...)
204-
end
205-
return a_dest
206-
end
207-
208-
# Only map the stored values of the inputs.
209-
function map_stored! end
210-
211-
@interface interface::AbstractArrayInterface function map_stored!(
212-
f, a_dest::AbstractArray, as::AbstractArray...
213-
)
214-
@interface interface map_indices!(eachstoredindex(as...), f, a_dest, as...)
215-
return a_dest
216-
end
217-
218-
# Only map all values, not just the stored ones.
219-
function map_all! end
220-
221-
@interface interface::AbstractArrayInterface function map_all!(
222-
f, a_dest::AbstractArray, as::AbstractArray...
223-
)
224-
@interface interface map_indices!(eachindex(as...), f, a_dest, as...)
225-
return a_dest
226-
end
227-
228114
using ArrayLayouts: ArrayLayouts, zero!
229115

230116
# `zero!` isn't defined in `Base`, but it is defined in `ArrayLayouts`
@@ -240,35 +126,6 @@ using ArrayLayouts: ArrayLayouts, zero!
240126
return @interface interface map_stored!(f, a, a)
241127
end
242128

243-
# Determines if a function preserves the stored values
244-
# of the destination sparse array.
245-
# The current code may be inefficient since it actually
246-
# accesses an unstored element, which in the case of a
247-
# sparse array of arrays can allocate an array.
248-
# Sparse arrays could be expected to define a cheap
249-
# unstored element allocator, for example
250-
# `get_prototypical_unstored(a::AbstractArray)`.
251-
function preserves_unstored(f, a_dest::AbstractArray, as::AbstractArray...)
252-
I = first(eachindex(as...))
253-
return iszero(f(map(a -> getunstoredindex(a, I), as)...))
254-
end
255-
256-
@interface interface::AbstractSparseArrayInterface function Base.map!(
257-
f, a_dest::AbstractArray, as::AbstractArray...
258-
)
259-
indices = if !preserves_unstored(f, a_dest, as...)
260-
eachindex(a_dest)
261-
elseif any(a -> a_dest !== a, as)
262-
as = map(a -> Base.unalias(a_dest, a), as)
263-
@interface interface zero!(a_dest)
264-
eachstoredindex(as...)
265-
else
266-
eachstoredindex(a_dest)
267-
end
268-
@interface interface map_indices!(indices, f, a_dest, as...)
269-
return a_dest
270-
end
271-
272129
# `f::typeof(norm)`, `op::typeof(max)` used by `norm`.
273130
function reduce_init(f, op, as...)
274131
# TODO: Generalize this.
@@ -304,6 +161,8 @@ struct SparseArrayStyle{N} <: AbstractSparseArrayStyle{N} end
304161

305162
SparseArrayStyle{M}(::Val{N}) where {M,N} = SparseArrayStyle{N}()
306163

164+
DerivableInterfaces.interface(::Type{<:AbstractSparseArrayStyle}) = SparseArrayInterface()
165+
307166
@interface ::AbstractSparseArrayInterface function Broadcast.BroadcastStyle(type::Type)
308167
return SparseArrayStyle{ndims(type)}()
309168
end

‎src/indexing.jl‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -230,3 +230,98 @@ end
230230
@interface ::AbstractArrayInterface storedvalues(A::AbstractArray) = values(A)
231231
@interface ::AbstractArrayInterface storedpairs(A::AbstractArray) = pairs(A)
232232
@interface ::AbstractArrayInterface storedlength(A::AbstractArray) = length(A)
233+
234+
# SparseArrayInterface implementations
235+
# ------------------------------------
236+
# canonical errors are moved to `isstored`, `getstoredindex` and `getunstoredindex`
237+
# so no errors at this level by defining both IndexLinear and IndexCartesian
238+
@interface ::AbstractSparseArrayInterface function Base.getindex(
239+
A::AbstractArray{<:Any,N}, I::Vararg{Int,N}
240+
) where {N}
241+
@_propagate_inbounds_meta
242+
@boundscheck checkbounds(A, I...) # generally isstored requires bounds checking
243+
return @inbounds isstored(A, I...) ? getstoredindex(A, I...) : getunstoredindex(A, I...)
244+
end
245+
@interface ::AbstractSparseArrayInterface function Base.getindex(A::AbstractArray, I::Int)
246+
@_propagate_inbounds_meta
247+
@boundscheck checkbounds(A, I)
248+
return @inbounds isstored(A, I) ? getstoredindex(A, I) : getunstoredindex(A, I)
249+
end
250+
# disambiguate vectors
251+
@interface ::AbstractSparseArrayInterface function Base.getindex(A::AbstractVector, I::Int)
252+
@_propagate_inbounds_meta
253+
@boundscheck checkbounds(A, I)
254+
return @inbounds isstored(A, I) ? getstoredindex(A, I) : getunstoredindex(A, I)
255+
end
256+
257+
@interface ::AbstractSparseArrayInterface function Base.setindex!(
258+
A::AbstractArray{<:Any,N}, v, I::Vararg{Int,N}
259+
) where {N}
260+
@_propagate_inbounds_meta
261+
@boundscheck checkbounds(A, I...)
262+
return @inbounds if isstored(A, I...)
263+
setstoredindex!(A, v, I...)
264+
else
265+
setunstoredindex!(A, v, I...)
266+
end
267+
end
268+
@interface ::AbstractSparseArrayInterface function Base.setindex!(A::AbstractArray, I::Int)
269+
@_propagate_inbounds_meta
270+
@boundscheck checkbounds(A, I)
271+
return @inbounds if isstored(A, I)
272+
setstoredindex!(A, v, I)
273+
else
274+
setunstoredindex!(A, v, I)
275+
end
276+
end
277+
# disambiguate vectors
278+
@interface ::AbstractSparseArrayInterface function Base.setindex!(A::AbstractVector, I::Int)
279+
@_propagate_inbounds_meta
280+
@boundscheck checkbounds(A, I)
281+
return @inbounds if isstored(A, I)
282+
setstoredindex!(A, v, I)
283+
else
284+
setunstoredindex!(A, v, I)
285+
end
286+
end
287+
288+
# required:
289+
@interface ::AbstractSparseArrayInterface eachstoredindex(
290+
style::IndexStyle, A::AbstractArray
291+
) = throw(MethodError(eachstoredindex, (style, A)))
292+
@interface ::AbstractSparseArrayInterface storedvalues(A::AbstractArray) =
293+
throw(MethodError(storedvalues, A))
294+
295+
# derived but may be specialized:
296+
@interface ::AbstractSparseArrayInterface function eachstoredindex(
297+
style::IndexStyle, A::AbstractArray, B::AbstractArray...
298+
)
299+
return union(map(Base.Fix1(eachstoredindex, style), (A, B...))...)
300+
end
301+
302+
@interface ::AbstractSparseArrayInterface storedlength(A::AbstractArray) =
303+
length(storedvalues(A))
304+
@interface ::AbstractSparseArrayInterface storedpairs(A::AbstractArray) =
305+
zip(eachstoredindex(A), storedvalues(A))
306+
307+
#=
308+
All sparse array interfaces are mapped through layout_getindex. (is this too opinionated?)
309+
310+
using ArrayLayouts getindex: this is a bit cumbersome because there already is a way to make that work focussed on types
311+
but here we want to focus on interfaces.
312+
eg: ArrayLayouts.@layoutgetindex ArrayType
313+
TODO: decide if we need the interface approach at all here
314+
=#
315+
for (Tr, Tc) in Iterators.product(
316+
Iterators.repeated((:Colon, :AbstractUnitRange, :AbstractVector, :Integer), 2)...
317+
)
318+
Tr === Tc === :Integer && continue
319+
@eval begin
320+
@interface ::AbstractSparseArrayInterface function Base.getindex(
321+
A::AbstractMatrix, kr::$Tr, jr::$Tc
322+
)
323+
Base.@inline # needed to make boundschecks work
324+
return ArrayLayouts.layout_getindex(A, kr, jr)
325+
end
326+
end
327+
end

‎src/sparsearrayinterface.jl‎

Lines changed: 0 additions & 26 deletions
This file was deleted.

0 commit comments

Comments
 (0)