Skip to content

Commit bd3236a

Browse files
andyferrisandyferris
authored
Fix linear indexing (#130)
* Fix linear indexing with colon and static arrays Fix #128 * setindex! tests
1 parent 314f114 commit bd3236a

File tree

2 files changed

+128
-11
lines changed

2 files changed

+128
-11
lines changed

‎src/indexing.jl

Lines changed: 105 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,9 @@
22
getindex(a::StaticArray, i::Int) = error("getindex(::$typeof(a), ::Int) is not defined.")
33
setindex!(a::StaticArray, value, i::Int) = error("setindex!(::$(typeof(a)), value, ::Int) is not defined.")
44

5-
######################
6-
## Scalar Indexing ##
7-
######################
5+
#######################################
6+
## Multidimensional scalar indexing ##
7+
#######################################
88

99
# Note: all indexing behavior defaults to dense, linear indexing
1010

@@ -48,8 +48,6 @@ end
4848
end
4949
end
5050

51-
52-
5351
#########################
5452
## Indexing utilities ##
5553
#########################
@@ -79,11 +77,109 @@ _ind(i::Int, ::Int, ::Type{Int}) = :(inds[$i])
7977
_ind(i::Int, j::Int, ::Type{<:StaticArray}) = :(inds[$i][$j])
8078
_ind(i::Int, j::Int, ::Type{Colon}) = j
8179

80+
################################
81+
## Non-scalar linear indexing ##
82+
################################
83+
84+
@inline function getindex(a::StaticArray, ::Colon)
85+
_getindex(a::StaticArray, Length(a), :)
86+
end
87+
88+
@generated function _getindex(a::StaticArray, ::Length{L}, ::Colon) where {L}
89+
exprs = [:(a[$i]) for i = 1:L]
90+
return quote
91+
@_inline_meta
92+
@inbounds return similar_type(a, Size(L))(tuple($(exprs...)))
93+
end
94+
end
95+
96+
@propagate_inbounds function getindex(a::StaticArray, inds::StaticArray{Int})
97+
_getindex(a, Length(inds), inds)
98+
end
99+
100+
@generated function _getindex(a::StaticArray, ::Length{L}, inds::StaticArray{Int}) where {L}
101+
exprs = [:(a[inds[$i]]) for i = 1:L]
102+
return quote
103+
@_propagate_inbounds_meta
104+
similar_type(a, Size(L))(tuple($(exprs...)))
105+
end
106+
end
107+
108+
@inline function setindex!(a::StaticArray, v, ::Colon)
109+
_setindex!(a::StaticArray, v, Length(a), :)
110+
return v
111+
end
82112

113+
@generated function _setindex!(a::StaticArray, v, ::Length{L}, ::Colon) where {L}
114+
exprs = [:(a[$i] = v) for i = 1:L]
115+
return quote
116+
@_inline_meta
117+
@inbounds $(Expr(:block, exprs...))
118+
end
119+
end
83120

84-
#####################
85-
## Array Indexing ##
86-
#####################
121+
@generated function _setindex!(a::StaticArray, v::AbstractArray, ::Length{L}, ::Colon) where {L}
122+
exprs = [:(a[$i] = v[$i]) for i = 1:L]
123+
return quote
124+
@_propagate_inbounds_meta
125+
if length(v) != L
126+
throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
127+
end
128+
@inbounds $(Expr(:block, exprs...))
129+
end
130+
end
131+
132+
@generated function _setindex!(a::StaticArray, v::StaticArray, ::Length{L}, ::Colon) where {L}
133+
exprs = [:(a[$i] = v[$i]) for i = 1:L]
134+
return quote
135+
@_propagate_inbounds_meta
136+
if Length(typeof(v)) != L
137+
throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
138+
end
139+
$(Expr(:block, exprs...))
140+
end
141+
end
142+
143+
@propagate_inbounds function setindex!(a::StaticArray, v, inds::StaticArray{Int})
144+
_setindex!(a, v, Length(inds), inds)
145+
return v
146+
end
147+
148+
@generated function _setindex!(a::StaticArray, v, ::Length{L}, inds::StaticArray{Int}) where {L}
149+
exprs = [:(a[inds[$i]] = v) for i = 1:L]
150+
return quote
151+
@_propagate_inbounds_meta
152+
similar_type(a, Size(L))(tuple($(exprs...)))
153+
end
154+
end
155+
156+
@generated function _setindex!(a::StaticArray, v::AbstractArray, ::Length{L}, inds::StaticArray{Int}) where {L}
157+
exprs = [:(a[$i] = v[$i]) for i = 1:L]
158+
return quote
159+
@_propagate_inbounds_meta
160+
if length(v) != L
161+
throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
162+
end
163+
$(Expr(:block, exprs...))
164+
end
165+
end
166+
167+
@generated function _setindex!(a::StaticArray, v::StaticArray, ::Length{L}, inds::StaticArray{Int}) where {L}
168+
exprs = [:(a[$i] = v[$i]) for i = 1:L]
169+
return quote
170+
@_propagate_inbounds_meta
171+
if Length(typeof(v)) != L
172+
throw(DimensionMismatch("tried to assign $(length(v))-element array to length-$L destination"))
173+
end
174+
$(Expr(:block, exprs...))
175+
end
176+
end
177+
178+
###########################################
179+
## Multidimensional non-scalar indexing ##
180+
###########################################
181+
182+
# getindex
87183

88184
@propagate_inbounds function getindex(a::StaticArray, inds::Union{Int, StaticArray{Int}, Colon}...)
89185
_getindex(a, index_sizes(Size(a), inds...), inds)
@@ -124,7 +220,6 @@ end
124220
_getindex(a, index_sizes(i1, i2, i3, i4, inds...), (i1, i2, i3, i4, inds...))
125221
end
126222

127-
128223
@generated function _getindex(a::AbstractArray, ind_sizes::Tuple{Vararg{Size}}, inds)
129224
newsize = out_index_size(ind_sizes.parameters...)
130225
linearsizes = linear_index_size(ind_sizes.parameters...)
@@ -161,7 +256,7 @@ end
161256
end
162257
end
163258

164-
259+
# setindex!
165260

166261
@propagate_inbounds function setindex!(a::StaticArray, value, inds::Union{Int, StaticArray{Int}, Colon}...)
167262
_setindex!(a, value, index_sizes(Size(a), inds...), inds)
@@ -264,7 +359,6 @@ end
264359
end
265360
end
266361

267-
268362
# setindex! from an array
269363
@generated function _setindex!(a::AbstractArray, v::AbstractArray, ind_sizes::Tuple{Vararg{Size}}, inds)
270364
linearsizes = linear_index_size(ind_sizes.parameters...)

‎test/indexing.jl

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,17 @@
99
@test (@inferred getindex(sv,:)) === sv
1010
end
1111

12+
@testset "Linear getindex() on SMatrix" begin
13+
sv = SVector{4}(4,5,6,7)
14+
sm = SMatrix{2,2}(4,5,6,7)
15+
16+
# SVector
17+
@test (@inferred getindex(sm, SVector(4,3,2,1))) === SVector((7,6,5,4))
18+
19+
# Colon
20+
@test (@inferred getindex(sm,:)) === sv
21+
end
22+
1223
@testset "Linear getindex()/setindex!() on MVector" begin
1324
vec = @SVector [4,5,6,7]
1425

@@ -21,6 +32,18 @@
2132
@test (mv[:] = vec; (@inferred getindex(mv, :))::SVector{4,Int} == SVector((4,5,6,7)))
2233
end
2334

35+
@testset "Linear getindex()/setindex!() on MMatrix" begin
36+
vec = @SVector [4,5,6,7]
37+
38+
# SVector
39+
mm = MMatrix{2,2,Int}()
40+
@test (mm[SVector(1,2,3,4)] = vec; (@inferred getindex(mm, SVector(4,3,2,1)))::SVector{4,Int} == SVector((7,6,5,4)))
41+
42+
# Colon
43+
mm = MMatrix{2,2,Int}()
44+
@test (mm[:] = vec; (@inferred getindex(mm, :))::SVector{4,Int} == SVector((4,5,6,7)))
45+
end
46+
2447
@testset "Linear getindex()/setindex!() with a SVector on an Array" begin
2548
v = [11,12,13]
2649

0 commit comments

Comments
 (0)