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

Improve isassigned implementation #49827

Merged
merged 5 commits into from
May 20, 2023
Merged
Show file tree
Hide file tree
Changes from 2 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
14 changes: 0 additions & 14 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -604,20 +604,6 @@ end
size_to_strides(s, d) = (s,)
size_to_strides(s) = ()


function isassigned(a::AbstractArray, i::Integer...)
try
a[i...]
true
catch e
if isa(e, BoundsError) || isa(e, UndefRefError)
return false
else
rethrow()
end
end
end

function isstored(A::AbstractArray{<:Any,N}, I::Vararg{Integer,N}) where {N}
@boundscheck checkbounds(A, I...)
return true
Expand Down
1 change: 1 addition & 0 deletions base/indices.jl
Original file line number Diff line number Diff line change
Expand Up @@ -504,6 +504,7 @@ promote_rule(a::Type{IdentityUnitRange{T1}}, b::Type{IdentityUnitRange{T2}}) whe
IndexStyle(::Type{<:LinearIndices}) = IndexLinear()
axes(iter::LinearIndices) = map(axes1, iter.indices)
size(iter::LinearIndices) = map(length, iter.indices)
isassigned(iter::LinearIndices, i::Int) = checkbounds(Bool, iter, i)
function getindex(iter::LinearIndices, i::Int)
@inline
@boundscheck checkbounds(iter, i)
Expand Down
30 changes: 29 additions & 1 deletion base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -344,6 +344,13 @@ module IteratorsMD
Base.axes(iter::CartesianIndices{N,R}) where {N,R} = map(Base.axes1, iter.indices)
Base.IndexStyle(::Type{CartesianIndices{N,R}}) where {N,R} = IndexCartesian()
Base.has_offset_axes(iter::CartesianIndices) = Base.has_offset_axes(iter.indices...)
@propagate_inbounds function isassigned(iter::CartesianIndices{N,R}, I::Vararg{Int, N}) where {N,R}
for i in 1:N
isassigned(iter.indices[i], I[i]) || return false
end
return true
end

# getindex for a 0D CartesianIndices is necessary for disambiguation
@propagate_inbounds function Base.getindex(iter::CartesianIndices{0,R}) where {R}
CartesianIndex()
Expand Down Expand Up @@ -1565,7 +1572,28 @@ end
end

isassigned(a::AbstractArray, i::CartesianIndex) = isassigned(a, Tuple(i)...)
isassigned(a::AbstractArray, i::Union{Integer, CartesianIndex}...) = isassigned(a, CartesianIndex(i))
function isassigned(A::AbstractArray, i::Union{Integer, CartesianIndex}...)
isa(i, Tuple{Vararg{Int}}) || return isassigned(A, CartesianIndex(i...))
@boundscheck checkbounds(Bool, A, i...) || return false
S = IndexStyle(A)
ninds = length(i)
if isa(S, IndexLinear) && ninds != 1
return @inbounds isassigned(A, _to_linear_index(A, i...))
elseif ndims(A) != ninds
return @inbounds isassigned(A, _to_subscript_indices(A, i...)...)
else
try
A[i...]
true
catch e
if isa(e, BoundsError) || isa(e, UndefRefError)
return false
else
rethrow()
end
end
end
end

## permutedims

Expand Down
6 changes: 6 additions & 0 deletions base/permuteddimsarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,12 @@ end
val
end

function Base.isassigned(A::PermutedDimsArray{T,N,perm,iperm}, I::Vararg{Int,N}) where {T,N,perm,iperm}
@boundscheck checkbounds(Bool, A, I...) || return false
@inbounds x = isassigned(A.parent, genperm(I, iperm)...)
x
end

@inline genperm(I::NTuple{N,Any}, perm::Dims{N}) where {N} = ntuple(d -> I[perm[d]], Val(N))
@inline genperm(I, perm::AbstractVector{Int}) = genperm(I, (perm...,))

Expand Down
2 changes: 2 additions & 0 deletions base/range.jl
Original file line number Diff line number Diff line change
Expand Up @@ -901,6 +901,8 @@ end

## indexing

isassigned(r::AbstractRange, i::Int) = firstindex(r) <= i <= lastindex(i)

_in_unit_range(v::UnitRange, val, i::Integer) = i > 0 && val <= v.stop && val >= v.start

function getindex(v::UnitRange{T}, i::Integer) where T
Expand Down
13 changes: 13 additions & 0 deletions base/reshapedarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,19 @@ end
offset_if_vec(i::Integer, axs::Tuple{<:AbstractUnitRange}) = i + first(axs[1]) - 1
offset_if_vec(i::Integer, axs::Tuple) = i

@inline function isassigned(A::ReshapedArrayLF, index::Int)
@boundscheck checkbounds(Bool, A, index) || return false
@inbounds ret = isassigned(parent(A), index)
ret
end
@inline function isassigned(A::ReshapedArray{T,N}, indices::Vararg{Int, N}) where {T,N}
@boundscheck checkbounds(Bool, A, indices...) || return false
axp = axes(A.parent)
i = offset_if_vec(_sub2ind(size(A), indices...), axp)
I = ind2sub_rs(axp, A.mi, i)
@inbounds isassigned(A.parent, I...)
end

@inline function getindex(A::ReshapedArrayLF, index::Int)
@boundscheck checkbounds(A, index)
@inbounds ret = parent(A)[index]
Expand Down
31 changes: 31 additions & 0 deletions base/subarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -352,6 +352,37 @@ function setindex!(V::FastContiguousSubArray{<:Any, 1}, x, i::Int)
V
end

function isassigned(V::SubArray{T,N}, I::Vararg{Int,N}) where {T,N}
@inline
@boundscheck checkbounds(Bool, V, I...) || return false
@inbounds r = isassigned(V.parent, reindex(V.indices, I)...)
r
end
function isassigned(V::FastSubArray, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i) || return false
@inbounds r = isassigned(V.parent, V.offset1 + V.stride1*i)
r
end
function isassigned(V::FastContiguousSubArray, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i) || return false
@inbounds r = isassigned(V.parent, V.offset1 + i)
r
end
function isassigned(V::FastSubArray{<:Any, 1}, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i) || return false
@inbounds r = isassigned(V.parent, V.offset1 + V.stride1*i)
r
end
function isassigned(V::FastContiguousSubArray{<:Any, 1}, i::Int)
@inline
@boundscheck checkbounds(Bool, V, i) || return false
@inbounds r = isassigned(V.parent, V.offset1 + i)
r
end

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

Expand Down
2 changes: 2 additions & 0 deletions stdlib/LinearAlgebra/src/adjtrans.jl
Original file line number Diff line number Diff line change
Expand Up @@ -335,6 +335,8 @@ axes(v::AdjOrTransAbsVec) = (Base.OneTo(1), axes(v.parent)...)
axes(A::AdjOrTransAbsMat) = reverse(axes(A.parent))
IndexStyle(::Type{<:AdjOrTransAbsVec}) = IndexLinear()
IndexStyle(::Type{<:AdjOrTransAbsMat}) = IndexCartesian()
@propagate_inbounds Base.isassigned(v::AdjOrTransAbsVec, i::Int) = isassigned(v.parent, i-1+first(axes(v.parent)[1]))
@propagate_inbounds Base.isassigned(v::AdjOrTransAbsMat, i::Int, j::Int) = isassigned(v.parent, j, i)
@propagate_inbounds getindex(v::AdjOrTransAbsVec{T}, i::Int) where {T} = wrapperop(v)(v.parent[i-1+first(axes(v.parent)[1])])::T
@propagate_inbounds getindex(A::AdjOrTransAbsMat{T}, i::Int, j::Int) where {T} = wrapperop(A)(A.parent[j, i])::T
@propagate_inbounds setindex!(v::AdjOrTransAbsVec, x, i::Int) = (setindex!(v.parent, wrapperop(v)(x), i-1+first(axes(v.parent)[1])); v)
Expand Down
13 changes: 13 additions & 0 deletions stdlib/LinearAlgebra/src/bidiag.jl
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,19 @@ function bidiagzero(A::Bidiagonal{<:AbstractMatrix}, i, j)
end
end

@inline function Base.isassigned(A::Bidiagonal, i::Int, j::Int)
@boundscheck checkbounds(Bool, A, i, j) || return false
if i == j
return @inbounds isassigned(A.dv, i)
elseif A.uplo == 'U' && (i == j - 1)
return @inbounds isassigned(A.ev, i)
elseif A.uplo == 'L' && (i == j + 1)
return @inbounds isassigned(A.ev, j)
else
return true
end
end

@inline function getindex(A::Bidiagonal{T}, i::Integer, j::Integer) where T
@boundscheck checkbounds(A, i, j)
if i == j
Expand Down
10 changes: 10 additions & 0 deletions stdlib/LinearAlgebra/src/diagonal.jl
Original file line number Diff line number Diff line change
Expand Up @@ -139,6 +139,16 @@ function size(D::Diagonal,d::Integer)
return d<=2 ? length(D.diag) : 1
end

@inline function Base.isassigned(D::Diagonal, i::Int, j::Int)
@boundscheck checkbounds(Bool, D, i, j) || return false
if i == j
@inbounds r = isassigned(D.diag, i)
else
r = true
end
r
end

@inline function getindex(D::Diagonal, i::Int, j::Int)
@boundscheck checkbounds(D, i, j)
if i == j
Expand Down
3 changes: 3 additions & 0 deletions stdlib/LinearAlgebra/src/hessenberg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,9 @@ function Matrix{T}(H::UpperHessenberg) where T
return triu!(copyto!(Matrix{T}(undef, m, n), H.data), -1)
end

Base.isassigned(H::UpperHessenberg, i::Int, j::Int) =
i <= j+1 ? isassigned(H.data, i, j) : true

getindex(H::UpperHessenberg{T}, i::Integer, j::Integer) where {T} =
i <= j+1 ? convert(T, H.data[i,j]) : zero(T)

Expand Down
9 changes: 9 additions & 0 deletions stdlib/LinearAlgebra/src/symmetric.jl
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,15 @@ const RealHermSymComplexSym{T<:Real,S} = Union{Hermitian{T,S}, Symmetric{T,S}, S

size(A::HermOrSym, d) = size(A.data, d)
size(A::HermOrSym) = size(A.data)
@inline function Base.isassigned(A::HermOrSym, i::Int, j::Int)
@boundscheck checkbounds(Bool, A, i, j) || return false
@inbounds if i == j || ((A.uplo == 'U') == (i < j))
return isassigned(A.data, i, j)
else
return isassigned(A.data, j, i)
end
end

@inline function getindex(A::Symmetric, i::Integer, j::Integer)
@boundscheck checkbounds(A, i, j)
@inbounds if i == j
Expand Down
9 changes: 9 additions & 0 deletions stdlib/LinearAlgebra/src/triangular.jl
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,15 @@ function full!(A::UnitUpperTriangular)
B
end

Base.isassigned(A::UnitLowerTriangular, i::Int, j::Int) =
i > j ? isassigned(A.data, i, j) : true
Base.isassigned(A::LowerTriangular, i::Int, j::Int) =
i >= j ? isassigned(A.data, i, j) : true
Base.isassigned(A::UnitUpperTriangular, i::Int, j::Int) =
i < j ? isassigned(A.data, i, j) : true
Base.isassigned(A::UpperTriangular, i::Int, j::Int) =
i <= j ? isassigned(A.data, i, j) : true

getindex(A::UnitLowerTriangular{T}, i::Integer, j::Integer) where {T} =
i > j ? A.data[i,j] : ifelse(i == j, oneunit(T), zero(T))
getindex(A::LowerTriangular, i::Integer, j::Integer) =
Expand Down
26 changes: 26 additions & 0 deletions stdlib/LinearAlgebra/src/tridiag.jl
Original file line number Diff line number Diff line change
Expand Up @@ -413,6 +413,19 @@ end
det(A::SymTridiagonal; shift::Number=false) = det_usmani(A.ev, A.dv, A.ev, shift)
logabsdet(A::SymTridiagonal; shift::Number=false) = logabsdet(ldlt(A; shift=shift))

@inline function Base.isassigned(A::SymTridiagonal, i::Int, j::Int)
@boundscheck checkbounds(Bool, A, i, j) || return false
if i == j
return @inbounds isassigned(A.dv, i)
elseif i == j + 1
return @inbounds isassigned(A.ev, j)
elseif i + 1 == j
return @inbounds isassigned(A.ev, i)
else
return true
end
end

@inline function getindex(A::SymTridiagonal{T}, i::Integer, j::Integer) where T
@boundscheck checkbounds(A, i, j)
if i == j
Expand Down Expand Up @@ -604,6 +617,19 @@ function diag(M::Tridiagonal{T}, n::Integer=0) where T
end
end

@inline function Base.isassigned(A::Tridiagonal, i::Int, j::Int)
@boundscheck checkbounds(A, i, j)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this ok? in all other isassigned implementations, this returns false if the bound check fails, but this throws an error?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right. That should be fixed

if i == j
return @inbounds isassigned(A.d, i)
elseif i == j + 1
return @inbounds isassigned(A.dl, j)
elseif i + 1 == j
return @inbounds isassigned(A.du, i)
else
return true
end
end

@inline function getindex(A::Tridiagonal{T}, i::Integer, j::Integer) where T
@boundscheck checkbounds(A, i, j)
if i == j
Expand Down