Teach selectdim that its slices are slices (#26655)

This allows selectdim to return more optimized views that can be considered IndexLinear in some cases and make it more closely match the behavior of view with `:`s. Fixes issue reported at https://discourse.julialang.org/t/selectdim-fails-to-recognize-indexlinear/10053.
JuliaLang · Apr 3, 2018 · 956811c · 956811c
1 parent 266d82a
commit 956811c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/base/abstractarraymath.jl b/base/abstractarraymath.jl
@@ -116,7 +116,7 @@ julia> selectdim(A, 2, 3)
  7
 ```
 """
-@inline selectdim(A::AbstractArray, d::Integer, i) = _selectdim(A, d, i, setindex(axes(A), i, d))
+@inline selectdim(A::AbstractArray, d::Integer, i) = _selectdim(A, d, i, setindex(map(Slice, axes(A)), i, d))
 @noinline function _selectdim(A, d, i, idxs)
     d >= 1 || throw(ArgumentError("dimension must be ≥ 1"))
     nd = ndims(A)

diff --git a/test/arrayops.jl b/test/arrayops.jl
@@ -1872,6 +1872,10 @@ end
         @test @inferred(f26009(A, 2:2)) == reshape(2:4:20, 1, :)
         @test @inferred(f26009(A, 2)) == 2:4:20
     end
+    A = reshape(1:24, 4, 3, 2)
+    @test IndexStyle(selectdim(A, 1, 1)) == IndexStyle(view(A, 1, :, :)) == IndexLinear()
+    @test IndexStyle(selectdim(A, 2, 1)) == IndexStyle(view(A, :, 1, :)) == IndexCartesian()
+    @test IndexStyle(selectdim(A, 3, 1)) == IndexStyle(view(A, :, :, 1)) == IndexLinear()
 end
 
 ###