Skip to content

Commit 4adc061

Browse files
timholytimholy
authored
Merge pull request #571 from JuliaArrays/teh/getindex
Clean up indexing and dispatch
2 parents 23f8210 + b3b1eeb commit 4adc061

File tree

11 files changed

+51
-114
lines changed

11 files changed

+51
-114
lines changed

‎src/MMatrix.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ end
5757
end
5858
end
5959

60-
@inline convert(::Type{MMatrix{S1,S2}}, a::StaticArray{<:Any, T}) where {S1,S2,T} = MMatrix{S1,S2,T}(Tuple(a))
60+
@inline convert(::Type{MMatrix{S1,S2}}, a::StaticArray{<:Tuple, T}) where {S1,S2,T} = MMatrix{S1,S2,T}(Tuple(a))
6161
@inline MMatrix(a::StaticMatrix) = MMatrix{size(typeof(a),1),size(typeof(a),2)}(Tuple(a))
6262

6363
# Simplified show for the type

‎src/SMatrix.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ end
5252
end
5353
end
5454

55-
@inline convert(::Type{SMatrix{S1,S2}}, a::StaticArray{<:Any, T}) where {S1,S2,T} = SMatrix{S1,S2,T}(Tuple(a))
55+
@inline convert(::Type{SMatrix{S1,S2}}, a::StaticArray{<:Tuple, T}) where {S1,S2,T} = SMatrix{S1,S2,T}(Tuple(a))
5656
@inline SMatrix(a::StaticMatrix) = SMatrix{size(typeof(a),1),size(typeof(a),2)}(Tuple(a))
5757

5858
# Simplified show for the type

‎src/StaticArrays.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ const StaticMatrixLike{T} = Union{
9191
Diagonal{T, <:StaticVector{<:Any, T}}
9292
}
9393
const StaticVecOrMatLike{T} = Union{StaticVector{<:Any, T}, StaticMatrixLike{T}}
94-
const StaticArrayLike{T} = Union{StaticVecOrMatLike{T}, StaticArray{<:Any, T}}
94+
const StaticArrayLike{T} = Union{StaticVecOrMatLike{T}, StaticArray{<:Tuple, T}}
9595

9696
const AbstractScalar{T} = AbstractArray{T, 0} # not exported, but useful none-the-less
9797
const StaticArrayNoEltype{S, N, T} = StaticArray{S, T, N}

‎src/arraymath.jl

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -119,9 +119,7 @@ end
119119
end
120120
end
121121

122-
# ambiguity with AbstractRNG and non-Float64... possibly an optimized form in Base?
123-
@inline rand!(rng::MersenneTwister, a::SA) where {SA <: StaticArray{<:Any, Float64}} = _rand!(rng, Size(SA), a)
124-
@inline rand!(rng::MersenneTwister, a::SA) where {SA <: StaticArray{<:Tuple, Float64, <:Any}} = _rand!(rng, Size(SA), a)
122+
@inline rand!(rng::MersenneTwister, a::SA) where {SA <: StaticArray{<:Tuple, Float64}} = _rand!(rng, Size(SA), a)
125123

126124
@inline randn!(rng::AbstractRNG, a::SA) where {SA <: StaticArray} = _randn!(rng, Size(SA), a)
127125
@generated function _randn!(rng::AbstractRNG, ::Size{s}, a::SA) where {s, SA <: StaticArray}

‎src/broadcast.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@ using Base.Broadcast: _bcsm
1010
# A constructor that changes the style parameter N (array dimension) is also required
1111
struct StaticArrayStyle{N} <: AbstractArrayStyle{N} end
1212
StaticArrayStyle{M}(::Val{N}) where {M,N} = StaticArrayStyle{N}()
13-
BroadcastStyle(::Type{<:StaticArray{<:Any, <:Any, N}}) where {N} = StaticArrayStyle{N}()
14-
BroadcastStyle(::Type{<:Transpose{<:Any, <:StaticArray{<:Any, <:Any, N}}}) where {N} = StaticArrayStyle{N}()
15-
BroadcastStyle(::Type{<:Adjoint{<:Any, <:StaticArray{<:Any, <:Any, N}}}) where {N} = StaticArrayStyle{N}()
13+
BroadcastStyle(::Type{<:StaticArray{<:Tuple, <:Any, N}}) where {N} = StaticArrayStyle{N}()
14+
BroadcastStyle(::Type{<:Transpose{<:Any, <:StaticArray{<:Tuple, <:Any, N}}}) where {N} = StaticArrayStyle{N}()
15+
BroadcastStyle(::Type{<:Adjoint{<:Any, <:StaticArray{<:Tuple, <:Any, N}}}) where {N} = StaticArrayStyle{N}()
1616
# Precedence rules
1717
BroadcastStyle(::StaticArrayStyle{M}, ::DefaultArrayStyle{N}) where {M,N} =
1818
DefaultArrayStyle(Val(max(M, N)))

‎src/deque.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ end
7676
# could also be justified to live in src/indexing.jl
7777
import Base.setindex
7878
@propagate_inbounds setindex(a::StaticArray, x, index::Int) = _setindex(Length(a), a, convert(eltype(typeof(a)), x), index)
79-
@generated function _setindex(::Length{L}, a::StaticArray{<:Any,T}, x::T, index::Int) where {L, T}
79+
@generated function _setindex(::Length{L}, a::StaticArray{<:Tuple,T}, x::T, index::Int) where {L, T}
8080
exprs = [:(ifelse($i == index, x, a[$i])) for i = 1:L]
8181
return quote
8282
@_propagate_inbounds_meta

‎src/indexing.jl

Lines changed: 20 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -182,46 +182,32 @@ end
182182
## Multidimensional non-scalar indexing ##
183183
###########################################
184184

185-
# getindex
186-
187-
@propagate_inbounds function getindex(a::StaticArray, inds::Union{Int, StaticArray{<:Any, Int}, Colon}...)
188-
_getindex(a, index_sizes(Size(a), inds...), inds)
189-
end
190-
191-
# Hard to describe "Union{Int, StaticArray{<:Any, Int}} with at least one StaticArray{<:Any, Int}"
192-
# Here we require the first StaticArray{<:Any, Int} to be within the first four dimensions
193-
@propagate_inbounds function getindex(a::AbstractArray, i1::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
194-
_getindex(a, index_sizes(i1, inds...), (i1, inds...))
195-
end
196-
197-
@propagate_inbounds function getindex(a::AbstractArray, i1::Int, i2::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
198-
_getindex(a, index_sizes(i1, i2, inds...), (i1, i2, inds...))
199-
end
185+
# To intercept `A[i1, ...]` where all `i` indexes have static sizes,
186+
# create a wrapper used to mark non-scalar indexing operations.
187+
# We insert this at a point in the dispatch hierarchy where we can intercept any
188+
# `typeof(A)` (specifically, including dynamic arrays) without triggering ambiguities.
200189

201-
@propagate_inbounds function getindex(a::AbstractArray, i1::Int, i2::Int, i3::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
202-
_getindex(a, index_sizes(i1, i2, i3, inds...), (i1, i2, i3, inds...))
190+
struct StaticIndexing{I}
191+
ind::I
203192
end
193+
unwrap(i::StaticIndexing) = i.ind
204194

205-
@propagate_inbounds function getindex(a::AbstractArray, i1::Int, i2::Int, i3::Int, i4::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
206-
_getindex(a, index_sizes(i1, i2, i3, i4, inds...), (i1, i2, i3, i4, inds...))
195+
function Base.to_indices(A, I::Tuple{Vararg{Union{Integer, CartesianIndex, StaticArray{<:Tuple,Int}}}})
196+
inds = to_indices(A, axes(A), I)
197+
return map(StaticIndexing, inds)
207198
end
208199

209-
# Disambuguity methods for the above
210-
@propagate_inbounds function getindex(a::StaticArray, i1::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
211-
_getindex(a, index_sizes(i1, inds...), (i1, inds...))
212-
end
200+
# getindex
213201

214-
@propagate_inbounds function getindex(a::StaticArray, i1::Int, i2::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
215-
_getindex(a, index_sizes(i1, i2, inds...), (i1, i2, inds...))
202+
@propagate_inbounds function getindex(a::StaticArray, inds::Union{Int, StaticArray{<:Tuple, Int}, Colon}...)
203+
_getindex(a, index_sizes(Size(a), inds...), inds)
216204
end
217205

218-
@propagate_inbounds function getindex(a::StaticArray, i1::Int, i2::Int, i3::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
219-
_getindex(a, index_sizes(i1, i2, i3, inds...), (i1, i2, i3, inds...))
206+
function Base._getindex(::IndexStyle, A::AbstractArray, i1::StaticIndexing, I::StaticIndexing...)
207+
inds = (unwrap(i1), map(unwrap, I)...)
208+
return StaticArrays._getindex(A, index_sizes(inds...), inds)
220209
end
221210

222-
@propagate_inbounds function getindex(a::StaticArray, i1::Int, i2::Int, i3::Int, i4::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
223-
_getindex(a, index_sizes(i1, i2, i3, i4, inds...), (i1, i2, i3, i4, inds...))
224-
end
225211

226212
@generated function _getindex(a::AbstractArray, ind_sizes::Tuple{Vararg{Size}}, inds)
227213
newsize = out_index_size(ind_sizes.parameters...)
@@ -261,68 +247,13 @@ end
261247

262248
# setindex!
263249

264-
@propagate_inbounds function setindex!(a::StaticArray, value, inds::Union{Int, StaticArray{<:Any, Int}, Colon}...)
250+
@propagate_inbounds function setindex!(a::StaticArray, value, inds::Union{Int, StaticArray{<:Tuple, Int}, Colon}...)
265251
_setindex!(a, value, index_sizes(Size(a), inds...), inds)
266252
end
267253

268-
# Hard to describe "Union{Int, StaticArray{<:Any, Int}} with at least one StaticArray{<:Any, Int}"
269-
# Here we require the first StaticArray{<:Any, Int} to be within the first four dimensions
270-
@propagate_inbounds function setindex!(a::AbstractArray, value, i1::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
271-
_setindex!(a, value, index_sizes(i1, inds...), (i1, inds...))
272-
end
273-
274-
@propagate_inbounds function setindex!(a::AbstractArray, value, i1::Int, i2::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
275-
_setindex!(a, value, index_sizes(i1, i2, inds...), (i1, i2, inds...))
276-
end
277-
278-
@propagate_inbounds function setindex!(a::AbstractArray, value, i1::Int, i2::Int, i3::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
279-
_setindex!(a, value, index_sizes(i1, i2, i3, inds...), (i1, i2, i3, inds...))
280-
end
281-
282-
@propagate_inbounds function setindex!(a::AbstractArray, value, i1::Int, i2::Int, i3::Int, i4::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
283-
_setindex!(a, value, index_sizes(i1, i2, i3, i4, inds...), (i1, i2, i3, i4, inds...))
284-
end
285-
286-
# Disambiguity methods for the above
287-
@propagate_inbounds function setindex!(a::StaticArray, value, i1::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
288-
_setindex!(a, value, index_sizes(i1, inds...), (i1, inds...))
289-
end
290-
291-
@propagate_inbounds function setindex!(a::StaticArray, value, i1::Int, i2::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
292-
_setindex!(a, value, index_sizes(i1, i2, inds...), (i1, i2, inds...))
293-
end
294-
295-
@propagate_inbounds function setindex!(a::StaticArray, value, i1::Int, i2::Int, i3::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
296-
_setindex!(a, value, index_sizes(i1, i2, i3, inds...), (i1, i2, i3, inds...))
297-
end
298-
299-
@propagate_inbounds function setindex!(a::StaticArray, value, i1::Int, i2::Int, i3::Int, i4::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
300-
_setindex!(a, value, index_sizes(i1, i2, i3, i4, inds...), (i1, i2, i3, i4, inds...))
301-
end
302-
303-
# disambiguities from Base
304-
@propagate_inbounds function setindex!(a::Array, value, i1::StaticVector{<:Any, Int})
305-
_setindex!(a, value, index_sizes(i1), (i1,))
306-
end
307-
308-
@propagate_inbounds function setindex!(a::Array, value::AbstractArray, i1::StaticVector{<:Any, Int})
309-
_setindex!(a, value, index_sizes(i1), (i1,))
310-
end
311-
312-
@propagate_inbounds function setindex!(a::Array, value::AbstractArray, i1::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
313-
_setindex!(a, value, index_sizes(i1, inds...), (i1, inds...))
314-
end
315-
316-
@propagate_inbounds function setindex!(a::Array, value::AbstractArray, i1::Int, i2::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
317-
_setindex!(a, value, index_sizes(i1, i2, inds...), (i1, i2, inds...))
318-
end
319-
320-
@propagate_inbounds function setindex!(a::Array, value::AbstractArray, i1::Int, i2::Int, i3::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
321-
_setindex!(a, value, index_sizes(i1, i2, i3, inds...), (i1, i2, i3, inds...))
322-
end
323-
324-
@propagate_inbounds function setindex!(a::Array, value::AbstractArray, i1::Int, i2::Int, i3::Int, i4::StaticArray{<:Any, Int}, inds::Union{Int, StaticArray{<:Any, Int}}...)
325-
_setindex!(a, value, index_sizes(i1, i2, i3, i4, inds...), (i1, i2, i3, i4, inds...))
254+
function Base._setindex!(::IndexStyle, a::AbstractArray, value, i1::StaticIndexing, I::StaticIndexing...)
255+
inds = (unwrap(i1), map(unwrap, I)...)
256+
return StaticArrays._setindex!(a, value, index_sizes(inds...), inds)
326257
end
327258

328259
# setindex! from a scalar

‎src/mapreduce.jl

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -176,23 +176,23 @@ end
176176
# with an initial value v0 = true and false.
177177
#
178178
# TODO: change to use Base.reduce_empty/Base.reduce_first
179-
@inline iszero(a::StaticArray{<:Any,T}) where {T} = reduce((x,y) -> x && iszero(y), a, init=true)
179+
@inline iszero(a::StaticArray{<:Tuple,T}) where {T} = reduce((x,y) -> x && iszero(y), a, init=true)
180180

181-
@inline sum(a::StaticArray{<:Any,T}; dims=:) where {T} = reduce(+, a; dims=dims)
182-
@inline sum(f, a::StaticArray{<:Any,T}; dims=:) where {T} = mapreduce(f, +, a; dims=dims)
183-
@inline sum(f::Union{Function, Type}, a::StaticArray{<:Any,T}; dims=:) where {T} = mapreduce(f, +, a; dims=dims) # avoid ambiguity
181+
@inline sum(a::StaticArray{<:Tuple,T}; dims=:) where {T} = reduce(+, a; dims=dims)
182+
@inline sum(f, a::StaticArray{<:Tuple,T}; dims=:) where {T} = mapreduce(f, +, a; dims=dims)
183+
@inline sum(f::Union{Function, Type}, a::StaticArray{<:Tuple,T}; dims=:) where {T} = mapreduce(f, +, a; dims=dims) # avoid ambiguity
184184

185-
@inline prod(a::StaticArray{<:Any,T}; dims=:) where {T} = reduce(*, a; dims=dims)
186-
@inline prod(f, a::StaticArray{<:Any,T}; dims=:) where {T} = mapreduce(f, *, a; dims=dims)
187-
@inline prod(f::Union{Function, Type}, a::StaticArray{<:Any,T}; dims=:) where {T} = mapreduce(f, *, a; dims=dims)
185+
@inline prod(a::StaticArray{<:Tuple,T}; dims=:) where {T} = reduce(*, a; dims=dims)
186+
@inline prod(f, a::StaticArray{<:Tuple,T}; dims=:) where {T} = mapreduce(f, *, a; dims=dims)
187+
@inline prod(f::Union{Function, Type}, a::StaticArray{<:Tuple,T}; dims=:) where {T} = mapreduce(f, *, a; dims=dims)
188188

189-
@inline count(a::StaticArray{<:Any,Bool}; dims=:) = reduce(+, a; dims=dims)
189+
@inline count(a::StaticArray{<:Tuple,Bool}; dims=:) = reduce(+, a; dims=dims)
190190
@inline count(f, a::StaticArray; dims=:) = mapreduce(x->f(x)::Bool, +, a; dims=dims)
191191

192-
@inline all(a::StaticArray{<:Any,Bool}; dims=:) = reduce(&, a; dims=dims, init=true) # non-branching versions
192+
@inline all(a::StaticArray{<:Tuple,Bool}; dims=:) = reduce(&, a; dims=dims, init=true) # non-branching versions
193193
@inline all(f::Function, a::StaticArray; dims=:) = mapreduce(x->f(x)::Bool, &, a; dims=dims, init=true)
194194

195-
@inline any(a::StaticArray{<:Any,Bool}; dims=:) = reduce(|, a; dims=dims, init=false) # (benchmarking needed)
195+
@inline any(a::StaticArray{<:Tuple,Bool}; dims=:) = reduce(|, a; dims=dims, init=false) # (benchmarking needed)
196196
@inline any(f::Function, a::StaticArray; dims=:) = mapreduce(x->f(x)::Bool, |, a; dims=dims, init=false) # (benchmarking needed)
197197

198198
_mean_denom(a, dims::Colon) = length(a)

‎src/matrix_multiply.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ end
193193
end
194194
end
195195

196-
@generated function partly_unrolled_multiply(::Size{sa}, ::Size{sb}, a::StaticMatrix{<:Any, <:Any, Ta}, b::StaticArray{<:Any, Tb}) where {sa, sb, Ta, Tb}
196+
@generated function partly_unrolled_multiply(::Size{sa}, ::Size{sb}, a::StaticMatrix{<:Any, <:Any, Ta}, b::StaticArray{<:Tuple, Tb}) where {sa, sb, Ta, Tb}
197197
if sa[2] != sb[1]
198198
throw(DimensionMismatch("Tried to multiply arrays of size $sa and $sb"))
199199
end

‎src/triangular.jl

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ end
193193
end
194194
end
195195

196-
@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::StaticArray{<:Any,TA}, B::UpperTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
196+
@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::StaticArray{<:Tuple,TA}, B::UpperTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
197197
m = sa[1]
198198
if length(sa) == 1
199199
n = 1
@@ -225,7 +225,7 @@ end
225225
end
226226
end
227227

228-
@generated function _A_mul_Bc(::Size{sa}, ::Size{sb}, A::StaticArray{<:Any,TA}, B::UpperTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
228+
@generated function _A_mul_Bc(::Size{sa}, ::Size{sb}, A::StaticArray{<:Tuple,TA}, B::UpperTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
229229
m = sa[1]
230230
if length(sa) == 1
231231
n = 1
@@ -284,7 +284,7 @@ end
284284
end
285285
end
286286

287-
@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::StaticArray{<:Any,TA}, B::LowerTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
287+
@generated function _A_mul_B(::Size{sa}, ::Size{sb}, A::StaticArray{<:Tuple,TA}, B::LowerTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
288288
m = sa[1]
289289
if length(sa) == 1
290290
n = 1
@@ -316,7 +316,7 @@ end
316316
end
317317
end
318318

319-
@generated function _A_mul_Bc(::Size{sa}, ::Size{sb}, A::StaticArray{<:Any,TA}, B::LowerTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
319+
@generated function _A_mul_Bc(::Size{sa}, ::Size{sb}, A::StaticArray{<:Tuple,TA}, B::LowerTriangular{TB,<:StaticMatrix}) where {sa,sb,TA,TB}
320320
m = sa[1]
321321
if length(sa) == 1
322322
n = 1

0 commit comments

Comments
 (0)