Skip to content

Commit

Permalink
check compatible sizes
Browse files Browse the repository at this point in the history
  • Loading branch information
jishnub committed Jul 14, 2021
1 parent 184d12e commit 9f91760
Show file tree
Hide file tree
Showing 3 changed files with 32 additions and 9 deletions.
27 changes: 18 additions & 9 deletions src/OffsetArrays.jl
Original file line number Diff line number Diff line change
Expand Up @@ -340,30 +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(...)))
# 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::AbstractVector, ::Tuple{OffsetAxis}) = A
_reshape(A, inds) = reshape(A, inds)
# 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), 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_nov(parent(A), inds)
Base.reshape(A::OffsetArray, inds::Dims) = _reshape_nov(parent(A), inds)
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::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)
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
Expand Down
6 changes: 6 additions & 0 deletions src/utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,12 @@ _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
_checksize(::Colon, s) = true

_toaxis(i::Integer) = Base.OneTo(i)
_toaxis(i) = i

Expand Down
8 changes: 8 additions & 0 deletions test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1794,6 +1794,14 @@ end
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)
# the following is broken if rp doesn't define its own reshape
# we may fix it here but that's perhaps too much type-piracy
@test_broken reshape(rp, length(rp))
end

@testset "permutedims" begin
Expand Down

0 comments on commit 9f91760

Please sign in to comment.