Unify reshape methods accepting Colon and OneTo

base/reshapedarray.jl

@@ -119,31 +119,39 @@ julia> reshape(1:6, 2, 3)
 """
 reshape
 
-reshape(parent::AbstractArray, dims::IntOrInd...) = reshape(parent, dims)
-reshape(parent::AbstractArray, shp::Tuple{Union{Integer,OneTo}, Vararg{Union{Integer,OneTo}}}) = reshape(parent, to_shape(shp))
-reshape(parent::AbstractArray, dims::Tuple{Integer, Vararg{Integer}}) = reshape(parent, map(Int, dims))
-reshape(parent::AbstractArray, dims::Dims)        = _reshape(parent, dims)
+# we collect the vararg indices and only define methods for tuples of indices
+reshape(parent::AbstractArray, dims::Union{Integer,Colon,AbstractUnitRange}...) = reshape(parent, dims)
+reshape(parent::AbstractArray, dims::Tuple{Vararg{Integer}}) = reshape(parent, map(Int, dims))
+# Usually, custom arrays would only need to specialize `reshape(A::MyArray, dims::Tuple{Vararg{Int}})`, and other
+# integer types passed as `dims` would be converted to `Int`s
+reshape(parent::AbstractArray, dims::Dims) = _reshape(parent, dims)
 
 # Allow missing dimensions with Colon():
-reshape(parent::AbstractVector, ::Colon) = parent
-reshape(parent::AbstractVector, ::Tuple{Colon}) = parent
-reshape(parent::AbstractArray, dims::Int...) = reshape(parent, dims)
-reshape(parent::AbstractArray, dims::Integer...) = reshape(parent, dims)
-reshape(parent::AbstractArray, dims::Union{Integer,Colon}...) = reshape(parent, dims)
-reshape(parent::AbstractArray, dims::Tuple{Vararg{Union{Integer,Colon}}}) = reshape(parent, _reshape_uncolon(parent, dims))
-
-@noinline throw1(dims) = throw(DimensionMismatch(LazyString("new dimensions ", dims,
+# Replace `OneTo`s by their lengths, and convert colons to sizes using _reshape_uncolon
+# We add a level of indirection to avoid method ambiguities in reshape
+reshape(parent::AbstractArray, dims::Tuple{Vararg{Union{Integer,Colon,OneTo}}}) = _reshape_maybecolon(parent, dims)
+_reshape_maybecolon(parent::AbstractVector, ::Tuple{Colon}) = parent
+# helper function analogous to to_shape, but this doesn't cast Integers to Int
+# this allows dispatching to a specialized reshape(M::MyArray, dims::Tuple{Vararg{Integer}})
+_to_shape(i::OneTo) = length(i)
+_to_shape(i::Union{Integer, Colon}) = i
+_to_shape(t::Tuple) = map(_to_shape, t)
+_reshape_maybecolon(parent::AbstractArray, dims::Tuple{Vararg{Union{Integer,Colon,OneTo}}}) = reshape(parent, _reshape_uncolon(parent, _to_shape(dims)))
+
+@noinline _reshape_throwcolon(dims) = throw(DimensionMismatch(LazyString("new dimensions ", dims,
         " may have at most one omitted dimension specified by `Colon()`")))
-@noinline throw2(lenA, dims) = throw(DimensionMismatch(string("array size ", lenA,
+@noinline _reshape_throwsize(lenA, dims) = throw(DimensionMismatch(LazyString("array size ", lenA,
     " must be divisible by the product of the new dimensions ", dims)))
 
-@inline function _reshape_uncolon(A, _dims::Tuple{Vararg{Union{Integer, Colon}}})
+_reshape_uncolon(parent::AbstractArray, dims::Tuple) =  _reshape_uncolon(length(parent), dims)
+_reshape_uncolon(len::Integer, ::Tuple{Colon}) = len
+@inline function _reshape_uncolon(len::Integer, _dims::Tuple{Vararg{Union{Integer, Colon}}})
     # promote the dims to `Int` at least
     dims = map(x -> x isa Colon ? x : promote_type(typeof(x), Int)(x), _dims)
+    dims isa Tuple{Vararg{Integer}} && return dims
     pre = _before_colon(dims...)
     post = _after_colon(dims...)
-    _any_colon(post...) && throw1(dims)
-    len = length(A)
+    _any_colon(post...) && _reshape_throwcolon(dims)
     _reshape_uncolon_computesize(len, dims, pre, post)
 end
 @inline function _reshape_uncolon_computesize(len::Int, dims, pre::Tuple{Vararg{Int}}, post::Tuple{Vararg{Int}})
@@ -167,18 +175,20 @@ end
     (pre..., sz, post...)
 end
 @inline function _reshape_uncolon_computesize_nonempty(len, dims, pr)
-    iszero(pr) && throw2(len, dims)
+    iszero(pr) && _reshape_throwsize(len, dims)
     (quo, rem) = divrem(len, pr)
-    iszero(rem) || throw2(len, dims)
+    iszero(rem) || _reshape_throwsize(len, dims)
     quo
 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 _before_colon() = ()
 @inline _after_colon(dim::Any, tail...) =  _after_colon(tail...)
 @inline _after_colon(dim::Colon, tail...) = tail
+@inline _after_colon() = ()
 
 reshape(parent::AbstractArray{T,N}, ndims::Val{N}) where {T,N} = parent
 function reshape(parent::AbstractArray, ndims::Val{N}) where N

test/abstractarray.jl

@@ -2205,6 +2205,13 @@ end
         @test b isa Matrix{Int}
         @test b.ref === a.ref
     end
+    C = reshape(CartesianIndices((2,2)), big(4))
+    @test axes(C, 1) == 1:4
+    @test C == CartesianIndex.([(1,1), (2,1), (1,2), (2,2)])
+end
+@testset "reshape with OneTo and Colon" begin
+    @test reshape(1:3, Base.OneTo(1), :)    == reshape(1:3, 1, 3)
+    @test reshape(1:3, Base.OneTo(1), :, 1) == reshape(1:3, 1, 3, 1)
 end
 @testset "AbstractArrayMath" begin
     @testset "IsReal" begin

test/offsetarray.jl

@@ -2,12 +2,23 @@
 
 isdefined(Main, :OffsetArrays) || @eval Main include("testhelpers/OffsetArrays.jl")
 using .Main.OffsetArrays
-import .Main.OffsetArrays: IdOffsetRange
+import .Main.OffsetArrays: IdOffsetRange, no_offset_view
 using Random
 using LinearAlgebra
 using Base: IdentityUnitRange
 using Test
 
+no_offset_axes(x, d) = no_offset_view(axes(x, d))
+no_offset_axes(x) = map(no_offset_view, axes(x))
+
+function same_value(r1, r2)
+    length(r1) == length(r2) || return false
+    for (v1, v2) in zip(r1, r2)
+        v1 == v2 || return false
+    end
+    return true
+end
+
 if !isdefined(@__MODULE__, :T24Linear)
     include("testhelpers/arrayindexingtypes.jl")
 end
@@ -846,6 +857,175 @@ end
     a = OffsetArray(4:5, 5:6)
     @test reshape(a, :) === a
     @test reshape(a, (:,)) === a
+
+    A0 = [1 3; 2 4]
+    A = OffsetArray(A0, (-1,2))
+
+    B = reshape(A0, -10:-9, 9:10)
+    @test isa(B, OffsetArray{Int,2})
+    @test parent(B) == A0
+    @test axes(B) == IdentityUnitRange.((-10:-9, 9:10))
+    B = reshape(A, -10:-9, 9:10)
+    @test isa(B, OffsetArray{Int,2})
+    @test axes(B) == IdentityUnitRange.((-10:-9, 9:10))
+    b = reshape(A, -7:-4)
+    @test axes(b) == (IdentityUnitRange(-7:-4),)
+    @test isa(parent(b), Vector{Int})
+    @test parent(b) == A0[:]
+    a = OffsetArray(rand(3,3,3), -1:1, 0:2, 3:5)
+    # Offset axes are required for reshape(::OffsetArray, ::Val) support
+    b = reshape(a, Val(2))
+    @test isa(b, OffsetArray{Float64,2})
+    @test axes(b) == IdentityUnitRange.((-1:1, 1:9))
+    b = reshape(a, Val(4))
+    @test isa(b, OffsetArray{Float64,4})
+    @test axes(b) == (axes(a)..., IdentityUnitRange(1:1))
+
+    @test reshape(OffsetArray(-1:0, -1:0), :, 1) == reshape(-1:0, 2, 1)
+    @test reshape(OffsetArray(-1:2, -1:2), -2:-1, :) == reshape(-1:2, -2:-1, 2)
+
+    @test reshape(OffsetArray(-1:0, -1:0), :) == OffsetArray(-1:0, -1:0)
+    @test reshape(A, :) == reshape(A0, :)
+
+    # reshape with one Colon for AbstractArrays
+    B = reshape(A0, -10:-9, :)
+    @test B isa OffsetArray{Int,2}
+    @test parent(B) == A0
+    @test no_offset_axes(B, 1) == -10:-9
+    @test axes(B, 2) == axes(A0, 2)
+
+    B = reshape(A0, -10:-9, 3:3, :)
+    @test B isa OffsetArray{Int,3}
+    @test same_value(A0, B)
+    @test no_offset_axes(B, 1) == -10:-9
+    @test no_offset_axes(B, 2) == 3:3
+    @test axes(B, 3) == 1:2
+
+    B = reshape(A0, -10:-9, 3:4, :)
+    @test B isa OffsetArray{Int,3}
+    @test same_value(A0, B)
+    @test no_offset_axes(B, 1) == -10:-9
+    @test no_offset_axes(B, 2) == 3:4
+    @test axes(B, 3) == 1:1
+
+    # pop the parent
+    B = reshape(A, size(A))
+    @test B == A0
+    B = reshape(A, (Base.OneTo(2), 2))
+    @test B == A0
+    B = reshape(A, (2,:))
+    @test B == A0
+
+    # julialang/julia #33614
+    A = OffsetArray(-1:0, (-2,))
+    @test reshape(A, :) == A
+    Arsc = reshape(A, :, 1)
+    Arss = reshape(A, 2, 1)
+    @test Arsc[1,1] == Arss[1,1] == -1
+    @test Arsc[2,1] == Arss[2,1] == 0
+    @test_throws BoundsError Arsc[0,1]
+    @test_throws BoundsError Arss[0,1]
+    A = OffsetArray([-1,0], (-2,))
+    Arsc = reshape(A, :, 1)
+    Arsc[1,1] = 5
+    @test first(A) == 5
+
+    Vec64  = zeros(6)
+    ind_a_64 = 3
+    ind_a_32 =Int32.(ind_a_64)
+    @test reshape(Vec64, ind_a_32, :) == reshape(Vec64, ind_a_64, :)
+
+    R = reshape(zeros(6), 2, :)
+    @test R isa Matrix
+    @test axes(R) == (1:2, 1:3)
+    R = reshape(zeros(6,1), 2, :)
+    @test R isa Matrix
+    @test axes(R) == (1:2, 1:3)
+    R = reshape(zeros(6), 2:3, :)
+    @test axes(R) == (2:3, 1:3)
+    R = reshape(zeros(6,1), 2:3, :)
+    @test axes(R) == (2:3, 1:3)
+
+    R = reshape(zeros(6), 1:2, :)
+    @test axes(R) == (1:2, 1:3)
+    R = reshape(zeros(6,1), 1:2, :)
+    @test axes(R) == (1:2, 1:3)
+
+    # 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
+
+    # reshape with one single colon becomes a `vec`
+    A = OffsetArray(rand(4, 4), -1, -1)
+    @test reshape(A, (:, )) == vec(A)
+    @test reshape(A, :) == vec(A)
+
+    A0 = [1 3; 2 4]
+    A = reshape(A0, 2:3, 4:5)
+    @test axes(A) == Base.IdentityUnitRange.((2:3, 4:5))
+    @test reshape(A, axes(parent(A))..., :) == reshape(A0, axes(A0)..., :)
+
+    B = reshape(A0, -10:-9, 9:10)
+    @test isa(B, OffsetArray{Int,2})
+    @test parent(B) == A0
+    @test axes(B) == Base.IdentityUnitRange.((-10:-9, 9:10))
+
+    @testset "BigInt reshape" begin
+        struct MyBigFill{T,N} <: AbstractArray{T,N}
+            val :: T
+            axes :: NTuple{N,Base.OneTo{BigInt}}
+        end
+        MyBigFill(val, sz::Tuple{}) = MyBigFill{typeof(val),0}(val, sz)
+        MyBigFill(val, sz::Tuple{Vararg{Integer}}) = MyBigFill(val, map(Base.OneTo{BigInt}, sz))
+        Base.size(M::MyBigFill) = map(length, M.axes)
+        Base.axes(M::MyBigFill) = M.axes
+        function Base.getindex(M::MyBigFill{<:Any,N}, ind::Vararg{Integer,N}) where {N}
+            checkbounds(M, ind...)
+            M.val
+        end
+        function Base.isassigned(M::MyBigFill{<:Any,N}, ind::Vararg{BigInt,N}) where {N}
+            checkbounds(M, ind...)
+            true
+        end
+        function Base.reshape(M::MyBigFill, ind::NTuple{N,BigInt}) where {N}
+            length(M) == prod(ind) || throw(ArgumentError("length mismatch in reshape"))
+            MyBigFill(M.val, ind)
+        end
+        Base.reshape(M::MyBigFill, ind::Tuple{}) = MyBigFill(M.val, ind)
+
+        M = MyBigFill(4, (2, 3))
+        O = OffsetArray(M)
+        @test vec(O) isa MyBigFill
+        @test vec(O) == vec(M)
+        @test reshape(O, :, big(2)) == reshape(M, :, big(2))
+        @test reshape(O, :, big(2)) isa MyBigFill
+        @test reshape(O, axes(M)) isa MyBigFill
+        @test reshape(O, axes(M)) == reshape(M, axes(M))
+        @test reshape(O, axes(M,1), :) isa MyBigFill
+        @test reshape(O, axes(M,1), :) == reshape(M, axes(M,1), :)
+        @test reshape(O, axes(M,1), big(1), :) isa MyBigFill
+        @test reshape(O, axes(M,1), big(1), :) == reshape(M, axes(M,1), big(1), :)
+
+        M = MyBigFill(4, (1,1))
+        O = OffsetArray(M)
+        @test reshape(O) isa MyBigFill
+        @test reshape(O) == reshape(M)
+    end
 end
 
 @testset "stack" begin
@@ -914,14 +1094,3 @@ end
     b = sum(a, dims=1)
     @test b[begin] == sum(r)
 end
-
-@testset "reshape" begin
-    A0 = [1 3; 2 4]
-    A = reshape(A0, 2:3, 4:5)
-    @test axes(A) == Base.IdentityUnitRange.((2:3, 4:5))
-
-    B = reshape(A0, -10:-9, 9:10)
-    @test isa(B, OffsetArray{Int,2})
-    @test parent(B) == A0
-    @test axes(B) == Base.IdentityUnitRange.((-10:-9, 9:10))
-end

test/testhelpers/OffsetArrays.jl

@@ -546,9 +546,6 @@ _similar_axes_or_length(A, T, ax::I, ::I) where {I} = similar(A, T, map(_indexle
 _similar_axes_or_length(AT, ax::I, ::I) where {I} = similar(AT, map(_indexlength, ax))
 
 # reshape accepts a single colon
-# this method is limited to AbstractUnitRange{<:Integer} to avoid method overwritten errors if Base defines the same,
-# see https://github.com/JuliaLang/julia/pull/56850
-Base.reshape(A::AbstractArray, inds::Union{Integer, Colon, AbstractUnitRange{<:Integer}}...) = reshape(A, inds)
 function Base.reshape(A::AbstractArray, inds::Tuple{Vararg{OffsetAxis}})
     AR = reshape(no_offset_view(A), map(_indexlength, inds))
     O = OffsetArray(AR, map(_offset, axes(AR), inds))