Merge pull request #569 from mbauman:mb/slices

Author: Chris Foster

src/SOneTo.jl

@@ -19,6 +19,11 @@ Base.axes(s::SOneTo) = (s,)
 Base.size(s::SOneTo) = (length(s),)
 Base.length(s::SOneTo{n}) where {n} = n
 
+# The axes of a Slice'd SOneTo use the SOneTo itself
+Base.axes(S::Base.Slice{<:SOneTo}) = (S.indices,)
+Base.unsafe_indices(S::Base.Slice{<:SOneTo}) = (S.indices,)
+Base.axes1(S::Base.Slice{<:SOneTo}) = S.indices
+
 @propagate_inbounds function Base.getindex(s::SOneTo, i::Int)
     @boundscheck checkbounds(s, i)
     return i

src/abstractarray.jl

@@ -107,6 +107,7 @@ similar(::Type{SA}) where {SA<:StaticArray} = similar(SA,eltype(SA))
 similar(::SA,::Type{T}) where {SA<:StaticArray,T} = similar(SA,T,Size(SA))
 similar(::Type{SA},::Type{T}) where {SA<:StaticArray,T} = similar(SA,T,Size(SA))
 
+# Cases where a Size is given as the dimensions
 similar(::A,s::Size{S}) where {A<:AbstractArray,S} = similar(A,eltype(A),s)
 similar(::Type{A},s::Size{S}) where {A<:AbstractArray,S} = similar(A,eltype(A),s)
 
@@ -119,22 +120,19 @@ similar(::Type{A},::Type{T},s::Size{S}) where {A<:AbstractArray,T,S} = mutable_s
 similar(::Type{SA},::Type{T},s::Size{S}) where {SA<:SizedArray,T,S} = sizedarray_similar_type(T,s,length_val(s))(undef)
 similar(::Type{A},::Type{T},s::Size{S}) where {A<:Array,T,S} = sizedarray_similar_type(T,s,length_val(s))(undef)
 
-# We should be able to deal with SOneTo axes
-similar(::A, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar(A, eltype(A), shape)
-similar(::Type{A}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar(A, eltype(A), shape)
-
-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)))
+# Support tuples of mixtures of `SOneTo`s alongside the normal `Integer` and `OneTo` options
+# by simply converting them to either a tuple of Ints or a Size, re-dispatching to either one
+# of the above methods (in the case of Size) or a base fallback (in the case of Ints).
+const HeterogeneousShape = Union{Integer, Base.OneTo, SOneTo}
 
-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))))
+similar(A::AbstractArray, ::Type{T}, shape::Tuple{HeterogeneousShape, Vararg{HeterogeneousShape}}) where {T} = similar(A, T, homogenize_shape(shape))
+similar(::Type{A}, shape::Tuple{HeterogeneousShape, Vararg{HeterogeneousShape}}) where {A<:AbstractArray} = similar(A, homogenize_shape(shape))
+# Use an Array for StaticArrays if we don't have a statically-known size
+similar(::Type{A}, shape::Tuple{Int, Vararg{Int}}) where {A<:StaticArray} = Array{eltype(A)}(undef, 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))
+homogenize_shape(::Tuple{}) = ()
+homogenize_shape(shape::Tuple{Vararg{SOneTo}}) = Size(map(last, shape))
+homogenize_shape(shape::Tuple{Vararg{HeterogeneousShape}}) = map(last, shape)
 
 
 @inline reshape(a::StaticArray, s::Size) = similar_type(a, s)(Tuple(a))

test/abstractarray.jl

@@ -68,12 +68,29 @@ 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}
+    @testset "similar and Base.Slice/IdentityUnitRange (issues #548, #556)" begin
+        v = @SVector [1,2,3]
+        m = @SMatrix [1 2 3; 4 5 6]
+        @test similar(v, Int, SOneTo(3)) isa MVector{3,Int}
+        @test similar(v, Int, SOneTo(3), SOneTo(4)) isa MMatrix{3,4,Int}
+        @test similar(v, Int, 3, SOneTo(4)) isa Matrix
+        @test similar(v, SOneTo(3)) isa MVector{3,Int}
+        @test similar(v, SOneTo(3), SOneTo(4)) isa MMatrix{3,4,Int}
+        @test similar(v, 3, SOneTo(4)) isa Matrix
+
+        @test m[:, 1:2] isa Matrix
+        @test m[:, [true, false, false]] isa Matrix
+        @test m[:, SOneTo(2)] isa MMatrix{2, 2, Int}
+        @test m[:, :] isa SMatrix{2, 3, Int}
+        @test m[:, 1] isa SVector{2, Int}
+        @test m[2, :] isa SVector{3, Int}
+
+        # Test case that failed in AstroLib.jl
+        r = @view(m[:, 2:3]) * @view(v[1:2])
+        @test r == m[:, 2:3] * v[1:2] == Array(m)[:, 2:3] * Array(v)[1:2]
     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]