Skip to content

Commit

Permalink
Improve isassigned implementation (JuliaLang#49827)
Browse files Browse the repository at this point in the history
Unless `isassigned` is called on `Array` with `Int`s, it uses a try
catch, which is notoriously slow. This PR provides changes the default
implementation of `isassigned` to coerce the indices provided to `Int`s
and converts to linear or cartesian indices, depending on the arrays
`IndexStyle`. This also overloads `isassigned` for many of the array
types defined in Base.

Fixes: JuliaLang#44720
  • Loading branch information
Tokazama authored and pull[bot] committed Feb 23, 2024
1 parent 64ee59a commit 8c9e12b
Show file tree
Hide file tree
Showing 14 changed files with 154 additions and 15 deletions.
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 (!isa(S, IndexLinear) && ninds != ndims(A))
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(r)

_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)
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

0 comments on commit 8c9e12b

Please sign in to comment.