Skip to content

Commit

Permalink
Fix unaliascopy(::SubArray) with indices of `Array{<:CartesianInd…
Browse files Browse the repository at this point in the history
…ex}` (#47779)

This PR fixes the bug caused by the trimming trick.
`Base.index_lengths` is not a proper tool to calculate the trimmed shape
as `indices` might consume more than 1 dim.
And we can avoid the unneeded "`repeat`" when we meet
`Array{CartesianIndex{0}}`

Test added.
Close #47644.
  • Loading branch information
N5N3 authored Dec 11, 2023
1 parent b71523e commit e39e77f
Show file tree
Hide file tree
Showing 2 changed files with 47 additions and 9 deletions.
39 changes: 32 additions & 7 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -110,16 +110,41 @@ unaliascopy(A::SubArray) = typeof(A)(unaliascopy(A.parent), map(unaliascopy, A.i

# When the parent is an Array we can trim the size down a bit. In the future this
# could possibly be extended to any mutable array.
function unaliascopy(V::SubArray{T,N,A,I,LD}) where {T,N,A<:Array,I<:Tuple{Vararg{Union{Real,AbstractRange,Array}}},LD}
dest = Array{T}(undef, index_lengths(V.indices...))
copyto!(dest, V)
function unaliascopy(V::SubArray{T,N,A,I,LD}) where {T,N,A<:Array,I<:Tuple{Vararg{Union{ScalarIndex,AbstractRange{<:ScalarIndex},Array{<:Union{ScalarIndex,AbstractCartesianIndex}}}}},LD}
dest = Array{T}(undef, _trimmedshape(V.indices...))
trimmedpind = _trimmedpind(V.indices...)
vdest = trimmedpind isa Tuple{Vararg{Union{Slice,Colon}}} ? dest : view(dest, trimmedpind...)
copyto!(vdest, view(V, _trimmedvind(V.indices...)...))
SubArray{T,N,A,I,LD}(dest, map(_trimmedindex, V.indices), 0, Int(LD))
end
# Get the proper trimmed shape
_trimmedshape(::ScalarIndex, rest...) = (1, _trimmedshape(rest...)...)
_trimmedshape(i::AbstractRange, rest...) = (maximum(i), _trimmedshape(rest...)...)
_trimmedshape(i::Union{UnitRange,StepRange,OneTo}, rest...) = (length(i), _trimmedshape(rest...)...)
_trimmedshape(i::AbstractArray{<:ScalarIndex}, rest...) = (length(i), _trimmedshape(rest...)...)
_trimmedshape(i::AbstractArray{<:AbstractCartesianIndex{0}}, rest...) = _trimmedshape(rest...)
_trimmedshape(i::AbstractArray{<:AbstractCartesianIndex{N}}, rest...) where {N} = (length(i), ntuple(Returns(1), Val(N - 1))..., _trimmedshape(rest...)...)
_trimmedshape() = ()
# We can avoid the repeation from `AbstractArray{CartesianIndex{0}}`
_trimmedpind(i, rest...) = (map(Returns(:), axes(i))..., _trimmedpind(rest...)...)
_trimmedpind(i::AbstractRange, rest...) = (i, _trimmedpind(rest...)...)
_trimmedpind(i::Union{UnitRange,StepRange,OneTo}, rest...) = ((:), _trimmedpind(rest...)...)
_trimmedpind(i::AbstractArray{<:AbstractCartesianIndex{0}}, rest...) = _trimmedpind(rest...)
_trimmedpind() = ()
_trimmedvind(i, rest...) = (map(Returns(:), axes(i))..., _trimmedvind(rest...)...)
_trimmedvind(i::AbstractArray{<:AbstractCartesianIndex{0}}, rest...) = (map(first, axes(i))..., _trimmedvind(rest...)...)
_trimmedvind() = ()
# Transform indices to be "dense"
_trimmedindex(i::Real) = oftype(i, 1)
_trimmedindex(i::AbstractUnitRange) = oftype(i, oneto(length(i)))
_trimmedindex(i::AbstractArray) = oftype(i, reshape(eachindex(IndexLinear(), i), axes(i)))

_trimmedindex(i::ScalarIndex) = oftype(i, 1)
_trimmedindex(i::AbstractRange) = i
_trimmedindex(i::Union{UnitRange,StepRange,OneTo}) = oftype(i, oneto(length(i)))
_trimmedindex(i::AbstractArray{<:ScalarIndex}) = oftype(i, reshape(eachindex(IndexLinear(), i), axes(i)))
_trimmedindex(i::AbstractArray{<:AbstractCartesianIndex{0}}) = oftype(i, copy(i))
function _trimmedindex(i::AbstractArray{<:AbstractCartesianIndex{N}}) where {N}
padding = ntuple(Returns(1), Val(N - 1))
ax1 = eachindex(IndexLinear(), i)
return oftype(i, reshape(CartesianIndices((ax1, padding...)), axes(i)))
end
## SubArray creation
# We always assume that the dimensionality of the parent matches the number of
# indices that end up getting passed to it, so we store the parent as a
Expand Down
17 changes: 15 additions & 2 deletions test/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -660,8 +660,21 @@ end
@testset "unaliascopy trimming; Issue #26263" begin
A = rand(5,5,5,5)
V = view(A, 2:5, :, 2:5, 1:2:5)
@test @inferred(Base.unaliascopy(V)) == V == A[2:5, :, 2:5, 1:2:5]
@test @inferred(sum(Base.unaliascopy(V))) sum(V) sum(A[2:5, :, 2:5, 1:2:5])
V′ = @inferred(Base.unaliascopy(V))
@test size(V′.parent) == size(V)
@test V′::typeof(V) == V == A[2:5, :, 2:5, 1:2:5]
@test @inferred(sum(V′)) sum(V) sum(A[2:5, :, 2:5, 1:2:5])
V = view(A, Base.IdentityUnitRange(2:4), :, Base.StepRangeLen(1,1,3), 1:2:5)
V′ = @inferred(Base.unaliascopy(V))
@test size(V.parent) != size(V′.parent)
@test V′ == V && V′ isa typeof(V)
i1 = collect(CartesianIndices((2:5)))
i2 = [CartesianIndex(), CartesianIndex()]
i3 = collect(CartesianIndices((2:5, 1:2:5)))
V = view(A, i1, 1:5, i2, i3)
@test @inferred(Base.unaliascopy(V))::typeof(V) == V == A[i1, 1:5, i2, i3]
V = view(A, i1, 1:5, i3, i2)
@test @inferred(Base.unaliascopy(V))::typeof(V) == V == A[i1, 1:5, i3, i2]
end

@testset "issue #27632" begin
Expand Down

0 comments on commit e39e77f

Please sign in to comment.