Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add _unsetindex! methods for SubArrays and CartesianIndexes #53383

Merged
merged 3 commits into from
Feb 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 14 additions & 1 deletion base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1472,7 +1472,20 @@ function _setindex!(::IndexCartesian, A::AbstractArray, v, I::Vararg{Int,M}) whe
r
end

_unsetindex!(A::AbstractArray, i::Integer) = _unsetindex!(A, to_index(i))
function _unsetindex!(A::AbstractArray, i::Integer...)
@_propagate_inbounds_meta
_unsetindex!(A, map(to_index, i)...)
end

jishnub marked this conversation as resolved.
Show resolved Hide resolved
function _unsetindex!(A::AbstractArray{T}, i::Int...) where T
# this provides a fallback method which is a no-op if the element is already unassigned
# such that copying into an uninitialized object generally always will work,
# even if the specific custom array type has not implemented `_unsetindex!`
@inline
@boundscheck checkbounds(A, i...)
allocatedinline(T) || @inbounds(!isassigned(A, i...)) || throw(MethodError(_unsetindex!, (A, i...)))
return A
end

"""
parent(A)
Expand Down
7 changes: 6 additions & 1 deletion base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -219,7 +219,12 @@ function _unsetindex!(A::Array, i::Int)
@inbounds _unsetindex!(GenericMemoryRef(A.ref, i))
return A
end

function _unsetindex!(A::Array, i::Int...)
@inline
@boundscheck checkbounds(A, i...)
@inbounds _unsetindex!(A, _to_linear_index(A, i...))
return A
end

# TODO: deprecate this (aligned_sizeof and/or elsize and/or sizeof(Some{T}) are more correct)
elsize(::Type{A}) where {T,A<:Array{T}} = aligned_sizeof(T)
Expand Down
6 changes: 6 additions & 0 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1610,6 +1610,12 @@ end
end
end

# _unsetindex
@propagate_inbounds function Base._unsetindex!(A::AbstractArray, i::CartesianIndex)
Base._unsetindex!(A, to_indices(A, (i,))...)
return A
end

## permutedims

## Permute array dims ##
Expand Down
19 changes: 19 additions & 0 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -410,6 +410,25 @@ function isassigned(V::FastSubArray{<:Any, 1}, i::Int)
r
end

function _unsetindex!(V::FastSubArray, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i)
@inbounds _unsetindex!(V.parent, _reindexlinear(V, i))
return V
end
function _unsetindex!(V::FastSubArray{<:Any,1}, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i)
@inbounds _unsetindex!(V.parent, _reindexlinear(V, i))
return V
end
function _unsetindex!(V::SubArray{T,N}, i::Vararg{Int,N}) where {T,N}
@inline
@boundscheck checkbounds(Bool, V, i...)
@inbounds _unsetindex!(V.parent, reindex(V.indices, i)...)
return V
end

IndexStyle(::Type{<:FastSubArray}) = IndexLinear()
IndexStyle(::Type{<:SubArray}) = IndexCartesian()

Expand Down
31 changes: 31 additions & 0 deletions test/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2019,3 +2019,34 @@ end
@test B == A
end
end

@testset "_unsetindex!" begin
struct MyMatrixUnsetIndexCartInds{T,A<:AbstractMatrix{T}} <: AbstractMatrix{T}
data :: A
end
Base.size(A::MyMatrixUnsetIndexCartInds) = size(A.data)
Base.getindex(M::MyMatrixUnsetIndexCartInds, i::Int, j::Int) = M.data[i,j]
Base.setindex!(M::MyMatrixUnsetIndexCartInds, v, i::Int, j::Int) = setindex!(M.data, v, i, j)
struct MyMatrixUnsetIndexLinInds{T,A<:AbstractMatrix{T}} <: AbstractMatrix{T}
data :: A
end
Base.size(A::MyMatrixUnsetIndexLinInds) = size(A.data)
Base.getindex(M::MyMatrixUnsetIndexLinInds, i::Int) = M.data[i]
Base.setindex!(M::MyMatrixUnsetIndexLinInds, v, i::Int) = setindex!(M.data, v, i)
Base.IndexStyle(::Type{<:MyMatrixUnsetIndexLinInds}) = IndexLinear()

function test_unsetindex(MT)
M = MT(ones(2,2))
M2 = MT(Matrix{BigFloat}(undef, 2,2))
copyto!(M, M2)
@test all(==(1), M)
M3 = MT(Matrix{BigFloat}(undef, 2,2))
for i in eachindex(M3)
@test !isassigned(M3, i)
end
M3 .= 1
@test_throws MethodError copyto!(M3, M2)
end
test_unsetindex(MyMatrixUnsetIndexCartInds)
test_unsetindex(MyMatrixUnsetIndexLinInds)
end
42 changes: 42 additions & 0 deletions test/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1055,4 +1055,46 @@ end
@test !isassigned(v, 1, 2) # inbounds but not assigned
@test !isassigned(v, 3, 3) # out-of-bounds
end

@testset "_unsetindex!" begin
function test_unsetindex(A, B)
copyto!(A, B)
for i in eachindex(A)
@test !isassigned(A, i)
end
end
@testset "dest IndexLinear, src IndexLinear" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(copy(p), ntuple(_->:, ndims(p))...)
B = view(similar(A), ntuple(_->:, ndims(p))...)
test_unsetindex(A, B)
test_unsetindex(p, B)
end
end

@testset "dest IndexLinear, src IndexCartesian" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(copy(p), ntuple(_->:, ndims(p))...)
B = view(similar(A), axes(A)...)
test_unsetindex(A, B)
test_unsetindex(p, B)
end
end

@testset "dest IndexCartesian, src IndexLinear" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(p, axes(p)...)
B = similar(A)
test_unsetindex(A, B)
end
end

@testset "dest IndexCartesian, src IndexCartesian" begin
for p in (fill(BigInt(2)), BigInt[1, 2], BigInt[1 2; 3 4])
A = view(p, axes(p)...)
B = view(similar(A), axes(A)...)
test_unsetindex(A, B)
end
end
end
end