Reshape with an offset parent array

src/OffsetArrays.jl

@@ -11,16 +11,16 @@ export OffsetArray, OffsetMatrix, OffsetVector
 
 const IIUR = IdentityUnitRange{<:AbstractUnitRange{<:Integer}}
 
-include("axes.jl")
-include("utils.jl")
-include("origin.jl")
-
 # Technically we know the length of CartesianIndices but we need to convert it first, so here we
 # don't put it in OffsetAxisKnownLength.
 const OffsetAxisKnownLength = Union{Integer, AbstractUnitRange}
 const OffsetAxis = Union{OffsetAxisKnownLength, Colon}
 const ArrayInitializer = Union{UndefInitializer, Missing, Nothing}
 
+include("axes.jl")
+include("utils.jl")
+include("origin.jl")
+
 ## OffsetArray
 """
     OffsetArray(A, indices...)
@@ -347,15 +347,24 @@ end
 
 # Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return
 # an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...)))
+# try to pass on the indices as received to the parent.
+# If the parent doesn't define an appropriate method, this will fall back to using the lengths
+# Short-circuit the case with matching indices to circumvent the Base restriction to 1-based indices
+_reshape(A::AbstractArray{<:Any,N}, ::NTuple{N,Union{Integer, AbstractUnitRange}}) where {N} = A
+_reshape(A::AbstractArray{<:Any,N}, inds::NTuple{N,OffsetAxis}) where {N} = (_colon_check(inds); A)
+_reshape(A, inds) = reshape(A, inds)
+_reshape_nov(A, inds) = no_offset_view(_reshape(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(parent(A), inds)
+Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(parent(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(parent(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,13 @@
 
 _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()
 
+_toaxis(i::Integer) = Base.OneTo(i)
+_toaxis(i::OffsetAxis) = i
+
 _strip_IdOffsetRange(r::IdOffsetRange) = parent(r)
 _strip_IdOffsetRange(r) = r
 
@@ -111,3 +113,20 @@ _filterreshapeinds(t::Tuple) = _filterreshapeinds(tail(t))
 _filterreshapeinds(t::Tuple{}) = t
 _popreshape(A::AbstractArray, ax::Tuple{Vararg{Base.OneTo}}, inds::Tuple{}) = no_offset_view(A)
 _popreshape(A::AbstractArray, ax, inds) = A
+
+# derived from Base._reshape_uncolon
+# https://github.com/JuliaLang/julia/blob/d29126a43ee289fc5ab8fcb3dc0e03f514605950/base/reshapedarray.jl#L119-L137
+@inline function _colon_check(dims)
+    @noinline throw1(dims) = throw(DimensionMismatch(string("new dimensions $(dims) ",
+        "may have at most one omitted dimension specified by `Colon()`")))
+    post = _after_colon(dims...)
+    _any_colon(post...) && throw1(dims)
+    nothing
+end
+@inline _any_colon() = false
+@inline _any_colon(dim::Colon, tail...) = true
+@inline _any_colon(dim::Any, tail...) = _any_colon(tail...)
+# @inline _before_colon(dim::Any, tail...) = (dim, _before_colon(tail...)...)
+# @inline _before_colon(dim::Colon, tail...) = ()
+@inline _after_colon(dim::Any, tail...) =  _after_colon(tail...)
+@inline _after_colon(dim::Colon, tail...) = tail

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

test/runtests.jl

@@ -1756,6 +1756,31 @@ 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, :)
 end
 
 @testset "permutedims" begin