Unwrap arrays for all_assigned

JuliaArrays · Oct 31, 2022 · 5b51fdc · 5b51fdc
1 parent bcccec1
commit 5b51fdc
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 5 deletions.
diff --git a/lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl b/lib/ArrayInterfaceCore/src/ArrayInterfaceCore.jl
@@ -124,8 +124,10 @@ end
 Returns the parent array that type `T` wraps.
 """
 parent_type(x) = parent_type(typeof(x))
-parent_type(::Type{Symmetric{T,S}}) where {T,S} = S
 parent_type(::Type{<:AbstractTriangular{T,S}}) where {T,S} = S
+parent_type(@nospecialize T::Type{<:Symmetric}) = fieldtype(T, :data)
+parent_type(@nospecialize T::Type{<:Hermitian}) = fieldtype(T, :data)
+parent_type(@nospecialize T::Type{<:UpperHessenberg}) = fieldtype(T, :data)
 parent_type(@nospecialize T::Type{<:PermutedDimsArray}) = fieldtype(T, :parent)
 parent_type(@nospecialize T::Type{<:Adjoint}) = fieldtype(T, :parent)
 parent_type(@nospecialize T::Type{<:Transpose}) = fieldtype(T, :parent)
@@ -333,13 +335,49 @@ false
 ```
 """
 function all_assigned(x)
-    for i in eachindex(x)
-        @inbounds(isassigned(x, i)) || return false
+    if is_forwarding_wrapper(x)
+        return all_assigned(buffer(x))
+    else
+        for i in eachindex(x)
+            @inbounds(isassigned(x, i)) || return false
+        end
+        return true
+    end
+end
+function all_assigned(x::SparseMatrixCSC)
+    all_assigned(x.colptr) && all_assigned(x.rowval) && all_assigned(x.nzval)
+end
+all_assigned(x::SparseVector) = all_assigned(x.nzind) && all_assigned(x.nzval)
+all_assigned(x::Union{PermutedDimsArray,Base.ReshapedArray,SubArray}) = all_assigned(parent(x))
+all_assigned(x::Union{Symmetric,Hermitian,UpperHessenberg}) = all_assigned(parent(x))
+all_assigned(x::Union{UpTri,LoTri,Adjoint,Transpose,Diagonal}) = all_assigned(parent(x))
+all_assigned(x::Union{SymTridiagonal,Bidiagonal}) = all_assigned(x.dv) && all_assigned(x.ev)
+function all_assigned(x::Tridiagonal)
+    all_assigned(x.dl) && all_assigned(x.d) && all_assigned(x.du) &&
+    (isdefined(x, :du2) ? all_assigned(x.du2) : true)
+end
+all_assigned(::Union{BitArray,Base.SimpleVector}) = true
+# all values of `Array` are assigned if composed of bits types
+function all_assigned(x::Array{T}) where {T}
+    if Base.isbitsunion(T)
+        return true
+    else
+        i = length(x)
+        while i > 0
+            ccall(:jl_array_isassigned, Cint, (Any, UInt), x, i) == 1 || return false
+            i -= 1
+        end
+        return true
     end
-    return true
 end
 # ranges shouldn't be undefined at any index so long as they aren't mutable
 all_assigned(x::AbstractRange) = !ismutable(typeof(x))
+@inline function all_assigned(x::Union{LinearIndices,CartesianIndices})
+    for inds in x.indices
+        all_assigned(inds) || return false
+    end
+    return true
+end
 
 """
     can_setindex(::Type{T}) -> Bool
@@ -1010,7 +1048,6 @@ See also: [`IndicesInfo`](@ref), [`parentdims`](@ref)
 """
 childdims(@nospecialize info::IndicesInfo) = getfield(_lower_info(info), 3)
 
-
 """
     instances_do_not_alias(::Type{T}) -> Bool
 

diff --git a/lib/ArrayInterfaceCore/test/runtests.jl b/lib/ArrayInterfaceCore/test/runtests.jl
@@ -70,6 +70,17 @@ end
 @test ArrayInterfaceCore.can_avx(ArrayInterfaceCore.can_avx) == false
 
 @testset "all_assigned" begin
+    x = 1:16
+    r = reshape(x, 4, 4)
+    s = Symmetric(r)
+    @test ArrayInterfaceCore.all_assigned(r)
+    @test ArrayInterfaceCore.all_assigned(Diagonal(x))
+    @test ArrayInterfaceCore.all_assigned(SymTridiagonal(s))
+    @test ArrayInterfaceCore.all_assigned(Tridiagonal(s))
+    @test ArrayInterfaceCore.all_assigned(LinearIndices((2, 2)))
+    @test ArrayInterfaceCore.all_assigned(sparse([1,2,3],[1,2,3],[1,2,3]))
+    @test ArrayInterfaceCore.all_assigned(sparsevec([1, 2, 0, 0, 3, 0]))
+    @test ArrayInterfaceCore.all_assigned(view(s, :, 2))
     @test ArrayInterfaceCore.all_assigned([1, 2])
     @test !ArrayInterfaceCore.all_assigned(Vector{Any}(undef, 10))
 end