Support Base.Slice axes (fixes #548)

JuliaArrays · Nov 19, 2018 · bb4aca8 · bb4aca8
1 parent 4f089af
commit bb4aca8
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/src/abstractarray.jl b/src/abstractarray.jl
@@ -126,6 +126,12 @@ similar(::Type{A}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray
 similar(::A,::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(shape)))
 similar(::Type{A},::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(shape)))
 
+const SOneToLike{n} = Union{SOneTo{n}, Base.Slice{SOneTo{n}}}
+deslice(ax::SOneTo) = ax
+deslice(ax::Base.Slice) = ax.indices
+similar(::A,::Type{T}, shape::Tuple{SOneToLike, Vararg{SOneToLike}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(deslice.(shape))))
+similar(::Type{A},::Type{T}, shape::Tuple{SOneToLike, Vararg{SOneToLike}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(deslice.(shape))))
+
 # Handle mixtures of SOneTo and other ranges (probably should make Base more robust here)
 similar(::Type{A}, shape::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {A<:AbstractArray} = similar(A, length.(shape)) # Jumps back to 2-argument form in Base
 similar(::Type{A},::Type{T}, shape::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {A<:AbstractArray,T} = similar(A, length.(shape))

diff --git a/test/abstractarray.jl b/test/abstractarray.jl
@@ -68,6 +68,12 @@ using StaticArrays, Test, LinearAlgebra
         @test isa(@inferred(similar(Matrix{Int}, Int, Size(2,2))), SizedArray{Tuple{2, 2}, Int, 2, 2})
     end
 
+    @testset "similar and Base.Slice (issue #548)" begin
+        v = @SVector [1]
+        b = similar(v, Int, Base.Slice(SOneTo(3)))
+        @test b isa MVector{3,Int}
+    end
+
     @testset "reshape" begin
         @test @inferred(reshape(SVector(1,2,3,4), Size(2,2))) === SMatrix{2,2}(1,2,3,4)
         @test @inferred(reshape([1,2,3,4], Size(2,2)))::SizedArray{Tuple{2,2},Int,2,1} == [1 3; 2 4]