Reshape with an offset parent array (#245)

* reshape with an offset parent array

* fix test on v1.0

* remove unnecesary functions

* add test

* remove indexoffset for Colon

* fix reshape bug with differing sizes

* check compatible sizes

* Add length for CustomRange in tests

* remove broken test for non-offsetarray type

* remove unnecessary method

Author: jishnub

src/OffsetArrays.jl

@@ -340,22 +340,39 @@ _similar_axes_or_length(AT, ax::I, ::I) where {I} = similar(AT, map(_indexlength
 # reshape accepts a single colon
 Base.reshape(A::AbstractArray, inds::OffsetAxis...) = reshape(A, inds)
 function Base.reshape(A::AbstractArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}})
-    AR = reshape(A, map(_indexlength, inds))
+    AR = reshape(no_offset_view(A), map(_indexlength, inds))
     O = OffsetArray(AR, map(_offset, axes(AR), inds))
     return _popreshape(O, axes(AR), _filterreshapeinds(inds))
 end
 
 # Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
 # an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
+# Short-circuit for AbstractVectors if the axes are compatible to get around the Base restriction
+# to 1-based vectors
+function _reshape(A::AbstractVector, inds::Tuple{OffsetAxis})
+    @noinline throw_dimerr(ind::Integer) = throw(
+        DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with length $ind"))
+    @noinline throw_dimerr(ind) = throw(
+        DimensionMismatch("parent has $(size(A,1)) elements, which is incompatible with indices $ind"))
+    _checksize(first(inds), size(A,1)) || throw_dimerr(first(inds))
+    A
+end
+_reshape(A, inds) = _reshape2(A, inds)
+_reshape2(A, inds) = reshape(A, inds)
+# avoid a stackoverflow by relegating to the parent if no_offset_view returns an offsetarray
+_reshape2(A::OffsetArray, inds) = reshape(parent(A), inds)
+_reshape_nov(A, inds) = _reshape(no_offset_view(A), inds)
+
 Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) =
-    OffsetArray(reshape(parent(A), map(_indexlength, inds)), map(_indexoffset, inds))
+    OffsetArray(_reshape(parent(A), inds), map(_toaxis, inds))
 # And for non-offset axes, we can just return a reshape of the parent directly
-Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = reshape(parent(A), inds)
-Base.reshape(A::OffsetArray, inds::Dims) = reshape(parent(A), inds)
-Base.reshape(A::OffsetArray, ::Colon) = reshape(parent(A), Colon())
+Base.reshape(A::OffsetArray, inds::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}) = _reshape_nov(A, inds)
+Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(A, inds)
 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)
+Base.reshape(A::OffsetVector, ::Tuple{Colon}) = A
+Base.reshape(A::OffsetArray, ::Colon) = reshape(A, (Colon(),))
+Base.reshape(A::OffsetArray, inds::Union{Int,Colon}...) = reshape(A, inds)
+Base.reshape(A::OffsetArray, inds::Tuple{Vararg{Union{Int,Colon}}}) = _reshape_nov(A, inds)
 
 # permutedims in Base does not preserve axes, and can not be fixed in a non-breaking way
 # This is a stopgap solution

src/utils.jl

@@ -2,11 +2,18 @@
 
 _indexoffset(r::AbstractRange) = first(r) - 1
 _indexoffset(i::Integer) = 0
-_indexoffset(i::Colon) = 0
 _indexlength(r::AbstractRange) = length(r)
 _indexlength(i::Integer) = Int(i)
 _indexlength(i::Colon) = Colon()
 
+# utility methods used in reshape
+# we don't use _indexlength in this to avoid converting the arguments to Int
+_checksize(ind::Integer, s) = ind == s
+_checksize(ind::AbstractUnitRange, s) = length(ind) == s
+
+_toaxis(i::Integer) = Base.OneTo(i)
+_toaxis(i) = i
+
 _strip_IdOffsetRange(r::IdOffsetRange) = parent(r)
 _strip_IdOffsetRange(r) = r
 

test/customranges.jl

@@ -65,6 +65,10 @@ for Z in [:ZeroBasedRange, :ZeroBasedUnitRange]
         @boundscheck checkbounds(A, r)
         OffsetArrays._indexedby(A[r.a], axes(r))
     end
+
+    @eval Base.reshape(z::$Z, inds::Tuple{}) = reshape(parent(z), inds)
+    @eval Base.reshape(z::$Z, inds::Tuple{Int, Vararg{Int}}) = reshape(parent(z), inds)
+    @eval Base.reshape(z::$Z, inds::Tuple{Union{Int, AbstractUnitRange{<:Integer}}, Vararg{Union{Int, AbstractUnitRange{<:Integer}}}}) = reshape(parent(z), inds)
 end
 
 # A basic range that does not have specialized vector indexing methods defined
@@ -75,6 +79,7 @@ struct CustomRange{T,A<:AbstractRange{T}} <: AbstractRange{T}
 end
 Base.parent(r::CustomRange) = r.a
 Base.size(r::CustomRange) = size(parent(r))
+Base.length(r::CustomRange) = length(parent(r))
 Base.axes(r::CustomRange) = axes(parent(r))
 Base.first(r::CustomRange) = first(parent(r))
 Base.last(r::CustomRange) = last(parent(r))

test/runtests.jl

@@ -1773,6 +1773,49 @@ end
     @test axes(R) == (1:2, 1:3)
     R = reshape(zeros(6,1), 1:2, :)
     @test axes(R) == (1:2, 1:3)
+
+    r = OffsetArray(ZeroBasedRange(3:4), 1);
+    @test reshape(r, 2) == 3:4
+    @test reshape(r, (2,)) == 3:4
+    @test reshape(r, :) == 3:4
+    @test reshape(r, (:,)) == 3:4
+
+    # getindex for a reshaped array that wraps an offset array is broken on 1.0
+    if VERSION >= v"1.1"
+        @test reshape(r, (2,:,4:4)) == OffsetArray(reshape(3:4, 2, 1, 1), 1:2, 1:1, 4:4)
+    end
+
+    # reshape works even if the parent doesn't have 1-based indices
+    # this works even if the parent doesn't support the reshape
+    r = OffsetArray(IdentityUnitRange(0:1), -1)
+    @test reshape(r, 2) == 0:1
+    @test reshape(r, (2,)) == 0:1
+    @test reshape(r, :) == OffsetArray(0:1, -1:0)
+    @test reshape(r, (:,)) == OffsetArray(0:1, -1:0)
+
+    @test reshape(ones(2:3, 4:5), (2, :)) == ones(2,2)
+
+    # more than one colon is not allowed
+    @test_throws Exception reshape(ones(3:4, 4:5, 1:2), :, :, 2)
+    @test_throws Exception reshape(ones(3:4, 4:5, 1:2), :, 2, :)
+
+    A = OffsetArray(rand(4, 4), -1, -1);
+    B = reshape(A, (2, :))
+    @test axes(B, 1) == 1:2
+    @test axes(B, 2) == 1:8
+
+    # some more exotic vector types
+    r = OffsetVector(CustomRange(ZeroBasedRange(0:2)), -2)
+    r2 = reshape(r, :)
+    @test r2 == r
+    r2 = reshape(r, 3)
+    @test axes(r2, 1) == 1:3
+    @test r2 == no_offset_view(r)
+    @test_throws Exception reshape(r, length(r) + 1)
+    @test_throws Exception reshape(r, 1:length(r) + 1)
+    rp = parent(r)
+    @test axes(reshape(rp, 4:6), 1) == 4:6
+    @test axes(reshape(r, (3,1))) == (1:3, 1:1)
 end
 
 @testset "permutedims" begin