Skip to content

Commit e3e63bd

Browse files
mbaumanc42f
authored andcommitted
Alternate approach to fix ##556
This restructures the similar method table along the lines that I suggested in #568 (comment). I also used a bit of internals magic to match how the Base `OneTo` operates -- this is a little fragile but I think it is a good way to replicate how `OneTo` works.
1 parent c6f2b5f commit e3e63bd

File tree

3 files changed

+38
-18
lines changed

3 files changed

+38
-18
lines changed

src/SOneTo.jl

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ Base.axes(s::SOneTo) = (s,)
1919
Base.size(s::SOneTo) = (length(s),)
2020
Base.length(s::SOneTo{n}) where {n} = n
2121

22+
# The axes of a Slice'd SOneTo use the SOneTo itself
23+
Base.axes(S::Base.Slice{<:SOneTo}) = (S.indices,)
24+
Base.unsafe_indices(S::Base.Slice{<:SOneTo}) = (S.indices,)
25+
Base.axes1(S::Base.Slice{<:SOneTo}) = S.indices
26+
2227
@propagate_inbounds function Base.getindex(s::SOneTo, i::Int)
2328
@boundscheck checkbounds(s, i)
2429
return i

src/abstractarray.jl

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ similar(::Type{SA}) where {SA<:StaticArray} = similar(SA,eltype(SA))
107107
similar(::SA,::Type{T}) where {SA<:StaticArray,T} = similar(SA,T,Size(SA))
108108
similar(::Type{SA},::Type{T}) where {SA<:StaticArray,T} = similar(SA,T,Size(SA))
109109

110+
# Cases where a Size is given as the dimensions
110111
similar(::A,s::Size{S}) where {A<:AbstractArray,S} = similar(A,eltype(A),s)
111112
similar(::Type{A},s::Size{S}) where {A<:AbstractArray,S} = similar(A,eltype(A),s)
112113

@@ -119,22 +120,19 @@ similar(::Type{A},::Type{T},s::Size{S}) where {A<:AbstractArray,T,S} = mutable_s
119120
similar(::Type{SA},::Type{T},s::Size{S}) where {SA<:SizedArray,T,S} = sizedarray_similar_type(T,s,length_val(s))(undef)
120121
similar(::Type{A},::Type{T},s::Size{S}) where {A<:Array,T,S} = sizedarray_similar_type(T,s,length_val(s))(undef)
121122

122-
# We should be able to deal with SOneTo axes
123-
similar(::A, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar(A, eltype(A), shape)
124-
similar(::Type{A}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray} = similar(A, eltype(A), shape)
125-
126-
similar(::A,::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(shape)))
127-
similar(::Type{A},::Type{T}, shape::Tuple{SOneTo, Vararg{SOneTo}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(shape)))
123+
# Support tuples of mixtures of `SOneTo`s alongside the normal `Integer` and `OneTo` options
124+
# by simply converting them to either a tuple of Ints or a Size, re-dispatching to either one
125+
# of the above methods (in the case of Size) or a base fallback (in the case of Ints).
126+
const HeterogeneousShape = Union{Integer, Base.OneTo, SOneTo}
128127

129-
const SOneToLike{n} = Union{SOneTo{n}, Base.Slice{SOneTo{n}}}
130-
deslice(ax::SOneTo) = ax
131-
deslice(ax::Base.Slice) = ax.indices
132-
similar(::A,::Type{T}, shape::Tuple{SOneToLike, Vararg{SOneToLike}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(deslice.(shape))))
133-
similar(::Type{A},::Type{T}, shape::Tuple{SOneToLike, Vararg{SOneToLike}}) where {A<:AbstractArray,T} = similar(A, T, Size(last.(deslice.(shape))))
128+
similar(A::AbstractArray, ::Type{T}, shape::Tuple{HeterogeneousShape, Vararg{HeterogeneousShape}}) where {T} = similar(A, T, homogenize_shape(shape))
129+
similar(::Type{A}, shape::Tuple{HeterogeneousShape, Vararg{HeterogeneousShape}}) where {A<:AbstractArray} = similar(A, homogenize_shape(shape))
130+
# Use an Array for StaticArrays if we don't have a statically-known size
131+
similar(::Type{A}, shape::Tuple{Int, Vararg{Int}}) where {A<:StaticArray} = Array{eltype(A)}(undef, shape)
134132

135-
# Handle mixtures of SOneTo and other ranges (probably should make Base more robust here)
136-
similar(::Type{A}, shape::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {A<:AbstractArray} = similar(A, length.(shape)) # Jumps back to 2-argument form in Base
137-
similar(::Type{A},::Type{T}, shape::Tuple{AbstractUnitRange, Vararg{AbstractUnitRange}}) where {A<:AbstractArray,T} = similar(A, length.(shape))
133+
homogenize_shape(::Tuple{}) = ()
134+
homogenize_shape(shape::Tuple{Vararg{SOneTo}}) = Size(map(last, shape))
135+
homogenize_shape(shape::Tuple{Vararg{HeterogeneousShape}}) = map(last, shape)
138136

139137

140138
@inline reshape(a::StaticArray, s::Size) = similar_type(a, s)(Tuple(a))

test/abstractarray.jl

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,12 +68,29 @@ using StaticArrays, Test, LinearAlgebra
6868
@test isa(@inferred(similar(Matrix{Int}, Int, Size(2,2))), SizedArray{Tuple{2, 2}, Int, 2, 2})
6969
end
7070

71-
@testset "similar and Base.Slice (issue #548)" begin
72-
v = @SVector [1]
73-
b = similar(v, Int, Base.Slice(SOneTo(3)))
74-
@test b isa MVector{3,Int}
71+
@testset "similar and Base.Slice/IdentityUnitRange (issues #548, #556)" begin
72+
v = @SVector [1,2,3]
73+
m = @SMatrix [1 2 3; 4 5 6]
74+
@test similar(v, Int, SOneTo(3)) isa MVector{3,Int}
75+
@test similar(v, Int, SOneTo(3), SOneTo(4)) isa MMatrix{3,4,Int}
76+
@test similar(v, Int, 3, SOneTo(4)) isa Matrix
77+
@test similar(v, SOneTo(3)) isa MVector{3,Int}
78+
@test similar(v, SOneTo(3), SOneTo(4)) isa MMatrix{3,4,Int}
79+
@test similar(v, 3, SOneTo(4)) isa Matrix
80+
81+
@test m[:, 1:2] isa Matrix
82+
@test m[:, [true, false, false]] isa Matrix
83+
@test m[:, SOneTo(2)] isa MMatrix{2, 2, Int}
84+
@test m[:, :] isa SMatrix{2, 3, Int}
85+
@test m[:, 1] isa SVector{2, Int}
86+
@test m[2, :] isa SVector{3, Int}
87+
88+
# Test case that failed in AstroLib.jl
89+
r = @view(m[:, 2:3]) * @view(v[1:2])
90+
@test r == m[:, 2:3] * v[1:2] == Array(m)[:, 2:3] * Array(v)[1:2]
7591
end
7692

93+
7794
@testset "reshape" begin
7895
@test @inferred(reshape(SVector(1,2,3,4), Size(2,2))) === SMatrix{2,2}(1,2,3,4)
7996
@test @inferred(reshape([1,2,3,4], Size(2,2)))::SizedArray{Tuple{2,2},Int,2,1} == [1 3; 2 4]

0 commit comments

Comments
 (0)