Merge pull request #17283 from JuliaLang/teh/fix_flatten

Loosen the signature on indices passed to _flatten. Fixes #17275.
JuliaLang · Jul 5, 2016 · b25739d · b25739d
2 parents 97ebeb9 + c68a2cd
commit b25739d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/base/multidimensional.jl b/base/multidimensional.jl
@@ -27,10 +27,11 @@ CartesianIndex{N}(index::NTuple{N,Integer}) = CartesianIndex{N}(index)
 (::Type{CartesianIndex{N}}){N}(index::Integer...) = CartesianIndex{N}(index)
 (::Type{CartesianIndex{N}}){N}() = CartesianIndex{N}(())
 # Un-nest passed CartesianIndexes
-CartesianIndex(index::Union{Integer, CartesianIndex}...) = CartesianIndex(_flatten((), index...))
-_flatten(out) = out
-@inline _flatten(out, i::Integer, I...)        = _flatten((out..., i), I...)
-@inline _flatten(out, i::CartesianIndex, I...) = _flatten((out..., i.I...), I...)
+CartesianIndex(index::Union{Integer, CartesianIndex}...) = CartesianIndex(flatten(index))
+Base.@pure flatten(I) = (_flatten(I...)...,)
+Base.@pure _flatten() = ()
+Base.@pure _flatten(i, I...)                 = (i, _flatten(I...)...)
+Base.@pure _flatten(i::CartesianIndex, I...) = (i.I..., _flatten(I...)...)
 CartesianIndex(index::Tuple{Vararg{Union{Integer, CartesianIndex}}}) = CartesianIndex(index...)
 
 # length
@@ -374,10 +375,10 @@ end
 end
 
 @propagate_inbounds function _getindex{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, I::Union{Real,AbstractArray,Colon,CartesianIndex}...)
-    getindex(A, IteratorsMD._flatten((), I...)...)
+    getindex(A, IteratorsMD.flatten(I)...)
 end
 @propagate_inbounds function _setindex!{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, v, I::Union{Real,AbstractArray,Colon,CartesianIndex}...)
-    setindex!(A, v, IteratorsMD._flatten((), I...)...)
+    setindex!(A, v, IteratorsMD.flatten(I)...)
 end
 
 ##

diff --git a/test/arrayops.jl b/test/arrayops.jl
@@ -1222,6 +1222,9 @@ a[1,CartesianIndex{2}(3,4)] = -2
 @test a[CartesianIndex{1}(2),3,CartesianIndex{1}(4)] == 44
 a[CartesianIndex{1}(2),3,CartesianIndex{1}(3)] = -3
 @test a[CartesianIndex{1}(2),3,CartesianIndex{1}(3)] == -3
+@test a[:, :, CartesianIndex((1,))] == a[:,:,1]
+@test a[CartesianIndex((1,)), [1,2], :] == a[1,[1,2],:]
+@test a[CartesianIndex((2,)), 3:4, :] == a[2,3:4,:]
 
 a = view(zeros(3, 4, 5), :, :, :)
 a[CartesianIndex{3}(2,3,3)] = -1