Skip to content

Commit

Permalink
Fix colon-reshaping of OffsetVector (#33890)
Browse files Browse the repository at this point in the history
* Fix colon-reshaping of OffsetVector
* reshape(::AbstractVector, ::Colon) is a no-op

(cherry picked from commit f80c3ee)
  • Loading branch information
timholy authored and KristofferC committed Dec 4, 2019
1 parent 7e4e17f commit ac555be
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 3 deletions.
12 changes: 9 additions & 3 deletions base/reshapedarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@ reshape(parent::AbstractArray, shp::Tuple{Union{Integer,OneTo}, Vararg{Union{Int
reshape(parent::AbstractArray, dims::Dims) = _reshape(parent, dims)

# Allow missing dimensions with Colon():
reshape(parent::AbstractVector, ::Colon) = parent
reshape(parent::AbstractArray, dims::Int...) = reshape(parent, dims)
reshape(parent::AbstractArray, dims::Union{Int,Colon}...) = reshape(parent, dims)
reshape(parent::AbstractArray, dims::Tuple{Vararg{Union{Int,Colon}}}) = _reshape(parent, _reshape_uncolon(parent, dims))
Expand Down Expand Up @@ -220,6 +221,8 @@ dataids(A::ReshapedArray) = dataids(A.parent)
d, r = divrem(ind, strds[1])
(_ind2sub_rs(front(ax), tail(strds), r)..., d + first(ax[end]))
end
offset_if_vec(i::Integer, axs::Tuple{<:AbstractUnitRange}) = i + first(axs[1]) - 1
offset_if_vec(i::Integer, axs::Tuple) = i

@inline function getindex(A::ReshapedArrayLF, index::Int)
@boundscheck checkbounds(A, index)
Expand All @@ -237,8 +240,9 @@ end
end

@inline function _unsafe_getindex(A::ReshapedArray{T,N}, indices::Vararg{Int,N}) where {T,N}
i = Base._sub2ind(size(A), indices...)
I = ind2sub_rs(axes(A.parent), A.mi, i)
axp = axes(A.parent)
i = offset_if_vec(Base._sub2ind(size(A), indices...), axp)
I = ind2sub_rs(axp, A.mi, i)
_unsafe_getindex_rs(parent(A), I)
end
@inline _unsafe_getindex_rs(A, i::Integer) = (@inbounds ret = A[i]; ret)
Expand All @@ -260,7 +264,9 @@ end
end

@inline function _unsafe_setindex!(A::ReshapedArray{T,N}, val, indices::Vararg{Int,N}) where {T,N}
@inbounds parent(A)[ind2sub_rs(axes(A.parent), A.mi, Base._sub2ind(size(A), indices...))...] = val
axp = axes(A.parent)
i = offset_if_vec(Base._sub2ind(size(A), indices...), axp)
@inbounds parent(A)[ind2sub_rs(axes(A.parent), A.mi, i)...] = val
val
end

Expand Down
13 changes: 13 additions & 0 deletions test/offsetarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -495,6 +495,19 @@ A = OffsetArray(rand(4,4), (-3,5))
@test vec(A) == reshape(A, :) == reshape(A, 16) == reshape(A, Val(1)) == A[:] == vec(A.parent)
A = OffsetArray(view(rand(4,4), 1:4, 4:-1:1), (-3,5))
@test vec(A) == reshape(A, :) == reshape(A, 16) == reshape(A, Val(1)) == A[:] == vec(A.parent)
# issue #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

# broadcast
a = [1]
Expand Down

0 comments on commit ac555be

Please sign in to comment.