diff --git a/src/OffsetArrays.jl b/src/OffsetArrays.jl index 61b918d6..094c2677 100644 --- a/src/OffsetArrays.jl +++ b/src/OffsetArrays.jl @@ -9,16 +9,16 @@ end export OffsetArray, OffsetMatrix, OffsetVector -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...) @@ -280,13 +280,15 @@ end 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)) - return OffsetArray(AR, map(_offset, axes(AR), inds)) + return OffsetArray(AR, _offset_reshape_uncolon(A, inds)) end # Reshaping OffsetArrays can "pop" the original OffsetArray wrapper and return # an OffsetArray(reshape(...)) instead of an OffsetArray(reshape(OffsetArray(...))) -Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) = - OffsetArray(reshape(parent(A), map(_indexlength, inds)), map(_indexoffset, inds)) +function Base.reshape(A::OffsetArray, inds::Tuple{OffsetAxis,Vararg{OffsetAxis}}) + AR = reshape(parent(A), map(_indexlength, inds)) + OffsetArray(AR, _offset_reshape_uncolon(A, inds)) +end # 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) diff --git a/src/utils.jl b/src/utils.jl index ba52fce7..ff691e7a 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -10,6 +10,37 @@ _indexlength(i::Colon) = Colon() _offset(axparent::AbstractUnitRange, ax::AbstractUnitRange) = first(ax) - first(axparent) _offset(axparent::AbstractUnitRange, ax::Integer) = 1 - first(axparent) +_offset_reshape_uncolon(::AbstractArray, I::Tuple{OffsetAxisKnownLength,Vararg{OffsetAxisKnownLength}}) = I +function _offset_reshape_uncolon(A::AbstractArray, I::Tuple{OffsetAxis,Vararg{OffsetAxis}}) + @noinline throw1(I) = throw(DimensionMismatch(string("new dimensions $(I) ", + "may have at most one omitted dimension specified by `Colon()`"))) + @noinline throw2(A, I) = throw(DimensionMismatch(string("array size $(length(A)) ", + "must be divisible by the product of the new dimensions $I"))) + + pre = _before_colon(I...) + post = _after_colon(I...) + _any_colon(post...) && throw1(dims) + nprev = isempty(pre) ? 1 : mapreduce(_length, *, pre) + npost = isempty(post) ? 1 : mapreduce(_length, *, post) + sz, remainder = divrem(length(A), nprev * npost) + remainder == 0 || throw2(A, dims) + + # preserve the offset information, the default offset beyond `ndims(A)` is 0 + n = length(pre) + r = Base.OneTo(Int(sz)) + Δ = n < ndims(A) ? _offset(axes(A, n+1), r) : 0 + (pre..., IdOffsetRange(r, -Δ), post...) +end +@inline _any_colon() = false +@inline _any_colon(r::Colon, tail...) = true +@inline _any_colon(r::Any, tail...) = _any_colon(tail...) +@inline _before_colon(r::Any, tail...) = (r, _before_colon(tail...)...) +@inline _before_colon(r::Colon, tail...) = () +@inline _after_colon(r::Any, tail...) = _after_colon(tail...) +@inline _after_colon(r::Colon, tail...) = tail +@inline _length(r::AbstractUnitRange) = length(r) +@inline _length(n::Int) = n + """ OffsetArrays.AxisConversionStyle(typeof(indices)) @@ -59,6 +90,7 @@ AxisConversionStyle(::Type) = SingleRange() AxisConversionStyle(::Type{<:CartesianIndices}) = TupleOfRanges() _convertTupleAbstractUnitRange(x) = _convertTupleAbstractUnitRange(AxisConversionStyle(typeof(x)), x) +_convertTupleAbstractUnitRange(::SingleRange, x::Int) = (Base.OneTo(x), ) _convertTupleAbstractUnitRange(::SingleRange, x) = (convert(AbstractUnitRange{Int}, x),) _convertTupleAbstractUnitRange(::TupleOfRanges, x) = convert(Tuple{Vararg{AbstractUnitRange{Int}}}, x) diff --git a/test/runtests.jl b/test/runtests.jl index 0d100b85..2f297067 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1478,6 +1478,124 @@ end Arsc = reshape(A, :, 1) Arsc[1,1] = 5 @test first(A) == 5 + + # reshape with one Colon for AbstractArrays + # If possible, offset is preserved for OffsetArrays + # we test the following cases: + # - (RangeOrInt, :) + # - (RangeOrInt, :, RangeOrInt) + # - (: RangeOrInt) + @testset "Colon" begin + A0 = [1 2 3; 4 5 6] + A = OffsetArray(A0, -1, -1) + + # (Range, :) + B = reshape(A0, -10:-9, :) + @test B isa OffsetArray{Int,2} + @test parent(B) === A0 + @test axes(B) == (-10:-9, 1:3) + B = reshape(A, -10:-9, :) + @test B isa OffsetArray{Int,2} + @test parent(B) === A0 + @test axes(B) == (-10:-9, 0:2) + + # (Range, Int, :) + B = reshape(A0, -10:-9, 1, :) + @test B isa OffsetArray{Int,3} + @test same_value(B, A0) + @test axes(B) == (-10:-9, 1:1, 1:3) + # The position of `:` exceeds `ndims(A)` thus the offset is assumed to be 0 there + B = reshape(A, -10:-9, 1, :) + @test B isa OffsetArray{Int,3} + @test same_value(B, A0) + @test axes(B) == (-10:-9, 1:1, 1:3) + + # (Range, Range, :) + B = reshape(A0, -10:-9, 3, :) + @test B isa OffsetArray{Int,3} + @test same_value(B, A0) + @test axes(B) == (-10:-9, 1:3, 1:1) + # The position of `:` exceeds `ndims(A)` thus the offset is assumed to be 0 there + B = reshape(A, -10:-9, -1:1, :) + @test B isa OffsetArray{Int,3} + @test same_value(B, A0) + @test axes(B) == (-10:-9, -1:1, 1:1) + + # (:, Int) uses Base implementaion + # skip it + + # (:, Range) + B = reshape(A0, :, -10:-8) + @test B isa OffsetArray{Int,2} + @test parent(B) === A0 + @test axes(B) == (1:2, -10:-8) + B = reshape(A, :, -10:-8) + @test B isa OffsetArray{Int,2} + @test parent(B) === A0 + @test axes(B) == (0:1, -10:-8) + + # (Range, :, Int) + B = reshape(A0, -10:-9, :, 1) + @test B isa OffsetArray{Int,3} + @test same_value(A0, B) + @test axes(B) == (-10:-9, 1:3, 1:1) + B = reshape(A0, -10:-9, :, 3) + @test B isa OffsetArray{Int,3} + @test same_value(A0, B) + @test axes(B) == (-10:-9, 1:1, 1:3) + + B = reshape(A, -10:-9, :, 1) + @test B isa OffsetArray{Int,3} + @test same_value(A, B) + @test axes(B) == (-10:-9, 0:2, 1:1) + B = reshape(A, -10:-9, :, 3) + @test B isa OffsetArray{Int,3} + @test same_value(A, B) + @test axes(B) == (-10:-9, 0:0, 1:3) + + # (Range, :, Range) + B = reshape(A0, -10:-9, :, -10:-10) + @test B isa OffsetArray{Int,3} + @test same_value(A0, B) + @test axes(B) == (-10:-9, 1:3, -10:-10) + B = reshape(A0, -10:-9, :, -10:-8) + @test B isa OffsetArray{Int,3} + @test same_value(A0, B) + @test axes(B) == (-10:-9, 1:1, -10:-8) + + B = reshape(A, -10:-9, :, -10:-10) + @test B isa OffsetArray{Int,3} + @test same_value(A, B) + @test axes(B) == (-10:-9, 0:2, -10:-10) + B = reshape(A, -10:-9, :, -10:-8) + @test B isa OffsetArray{Int,3} + @test same_value(A, B) + @test axes(B) == (-10:-9, 0:0, -10:-8) + + # (Int, :, Range) + B = reshape(A0, 1, :, -10:-9) + @test B isa OffsetArray{Int,3} + @test same_value(A0, B) + @test axes(B) == (1:1, 1:3, -10:-9) + B = reshape(A0, 3, :, -10:-9) + @test B isa OffsetArray{Int,3} + @test same_value(A0, B) + @test axes(B) == (1:3, 1:1, -10:-9) + + B = reshape(A, 1, :, -10:-9) + @test B isa OffsetArray{Int,3} + @test same_value(A, B) + @test axes(B) == (1:1, 0:2, -10:-9) + B = reshape(A, 3, :, -10:-9) + @test B isa OffsetArray{Int,3} + @test same_value(A, B) + @test axes(B) == (1:3, 0:0, -10:-9) + + @test_throws DimensionMismatch reshape(A0, 1:2, :, :) + @test_throws DimensionMismatch reshape(A0, 1:2, 2, :) + @test_throws DimensionMismatch reshape(A, 1:2, :, :) + @test_throws DimensionMismatch reshape(A, 1:2, 2, :) + end end @testset "Indexing with OffsetArray axes" begin