similar with IdOffsetRanges may use the parent axes types (#213)

* similar with IdOffsetRange

* Add test for similar(A, T, axes)

* allow a mix of axis types

* update comment

* update similar(::Type, axes)

* Add tests for mixed axis types

* fix comment

* rename functions

* rename helper function

Author: jishnub

src/OffsetArrays.jl

@@ -271,10 +271,34 @@ end
 
 Base.similar(A::OffsetArray, ::Type{T}, dims::Dims) where T =
     similar(parent(A), T, dims)
-function Base.similar(A::AbstractArray, ::Type{T}, inds::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T
-    B = similar(A, T, map(_indexlength, inds))
-    return OffsetArray(B, map(_offset, axes(B), inds))
+function Base.similar(A::AbstractArray, ::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where T
+    # strip IdOffsetRanges to extract the parent range and use it to generate the array
+    new_shape = map(_strip_IdOffsetRange, shape)
+    # route through _similar_axes_or_length to avoid a stack overflow if map(_strip_IdOffsetRange, shape) === shape
+    # This tries to use new_shape directly in similar if similar(A, T, ::typeof(new_shape)) is defined
+    # If this fails, it calls similar(A, T, map(_indexlength, new_shape)) to use the size along each axis
+    # to generate the new array
+    P = _similar_axes_or_length(A, T, new_shape, shape)
+    return OffsetArray(P, map(_offset, axes(P), shape))
 end
+function Base.similar(::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where {T<:AbstractArray}
+    new_shape = map(_strip_IdOffsetRange, shape)
+    P = _similar_axes_or_length(T, new_shape, shape)
+    OffsetArray(P, map(_offset, axes(P), shape))
+end
+# Try to use the axes to generate the parent array type
+# This is useful if the axes have special meanings, such as with static arrays
+# This method is hit if at least one axis provided to similar(A, T, axes) is an IdOffsetRange
+# For example this is hit when similar(A::OffsetArray) is called,
+# which expands to similar(A, eltype(A), axes(A))
+_similar_axes_or_length(A, T, ax, ::Any) = similar(A, T, ax)
+_similar_axes_or_length(AT, ax, ::Any) = similar(AT, ax)
+# Handle the general case by resorting to lengths along each axis
+# This is hit if none of the axes provided to similar(A, T, axes) are IdOffsetRanges,
+# and if similar(A, T, axes::AX) is not defined for the type AX.
+# In this case the best that we can do is to create a mutable array of the correct size
+_similar_axes_or_length(A, T, ax::I, ::I) where {I} = similar(A, T, map(_indexlength, ax))
+_similar_axes_or_length(AT, ax::I, ::I) where {I} = similar(AT, map(_indexlength, ax))
 
 # reshape accepts a single colon
 Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
@@ -295,11 +319,6 @@ Base.reshape(A::OffsetVector, ::Colon) = A
 Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(parent(A), inds)
 Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = reshape(parent(A), inds)
 
-function Base.similar(::Type{T}, shape::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) where {T<:AbstractArray}
-    P = T(undef, map(_indexlength, shape))
-    OffsetArray(P, map(_offset, axes(P), shape))
-end
-
 Base.fill(v, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {N} =
     fill!(similar(Array{typeof(v)}, inds), v)
 Base.zeros(::Type{T}, inds::NTuple{N, Union{Integer, AbstractUnitRange}}) where {T, N} =

src/utils.jl

@@ -7,6 +7,9 @@ _indexlength(r::AbstractRange) = length(r)
 _indexlength(i::Integer) = i
 _indexlength(i::Colon) = Colon()
 
+_strip_IdOffsetRange(r::IdOffsetRange) = parent(r)
+_strip_IdOffsetRange(r) = r
+
 _offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent)
 _offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent)
 

test/runtests.jl

@@ -1397,6 +1397,48 @@ end
     @test isa(B, OffsetArray{Int, 2})
     @test axes(B) == (0:0, 1:3)
 
+    s = @SVector[i for i in 1:10]
+    so = OffsetArray(s, 4);
+    @test typeof(parent(similar(so))) == typeof(similar(s))
+    @test typeof(parent(similar(so, eltype(so), axes(so)))) == typeof(similar(s))
+
+    s = SArray{Tuple{2,2,2},Int,3,8}((1,2,3,4,5,6,7,8))
+    so = OffsetArray(s, 0, 0, 0)
+    A = similar(so)
+    @test A isa OffsetArray
+    @test parent(A) isa StaticArray
+    for ax in Any[(axes(s,1), axes(so)[2:3]...), (axes(so,1), axes(s)[2:3]...)]
+        A = similar(so, Int, ax)
+        @test A isa OffsetArray
+        @test parent(A) isa StaticArray
+        @test axes(A) == ax
+    end
+
+    # check with an unseen axis type
+    A = similar(ones(1), Int, ZeroBasedUnitRange(3:4), 4:5)
+    @test eltype(A) === Int
+    @test axes(A) == (3:4, 4:5)
+    A = similar(ones(1), Int, IdOffsetRange(ZeroBasedUnitRange(3:4), 2), 4:5)
+    @test eltype(A) === Int
+    @test axes(A) == (5:6, 4:5)
+    A = similar(ones(1), Int, IdOffsetRange(ZeroBasedUnitRange(3:4), 2), 4)
+    @test eltype(A) === Int
+    @test axes(A) == (5:6, 1:4)
+
+    # test for similar(::Type, ax)
+    indsoffset = (IdOffsetRange(SOneTo(2), 2),)
+    A = similar(Array{Int}, indsoffset)
+    @test parent(A) isa StaticArray
+    @test axes(A) == (3:4,)
+    A = similar(Array{Int}, (indsoffset..., SOneTo(3)))
+    @test parent(A) isa StaticArray
+    @test axes(A) == (3:4, 1:3)
+
+    s = SArray{Tuple{2,2},Int,2,4}((1,2,3,4));
+    so = OffsetArray(s, 2, 2);
+    so2 = so .+ 1;
+    @test parent(so2) isa StaticArray
+
     @test_throws MethodError similar(A, (:,))
     @test_throws MethodError similar(A, (: ,:))
     @test_throws MethodError similar(A, (: ,2))