Add known_length(::Type{<:SubArray})

Author: Tokazama

src/ArrayInterface.jl

@@ -35,8 +35,19 @@ Otherwise, return `nothing`.
 known_length(x) = known_length(typeof(x))
 known_length(::Type{<:NTuple{N,<:Any}}) where {N} = N
 known_length(::Type{<:NamedTuple{L}}) where {L} = length(L)
-known_length(::Type{T}) where {T<:Base.Slice} = known_length(parent_type(T))
-known_length(::Type{T}) where {T} = nothing
+known_length(::Type{<:Number}) = 1
+function known_length(::Type{T}) where {T}
+    if parent_type(T) <: T
+        return nothing
+    else
+        return known_length(parent_type(T))
+    end
+end
+@inline function known_length(::Type{<:SubArray{T,N,P,I}}) where {T,N,P,I}
+    return _known_length(ntuple(i -> known_length(I.parameters[i]), Val(N)))
+end
+_known_length(x::Tuple{Vararg{<:Union{Int,Nothing}}}) = nothing
+_known_length(x::Tuple{Vararg{Int}}) = prod(x)
 
 """
     can_change_size(::Type{T}) -> Bool

test/runtests.jl

@@ -205,6 +205,10 @@ end
     @test ArrayInterface.known_length((1,)) == 1
     @test ArrayInterface.known_length((a=1,b=2)) == 2
     @test ArrayInterface.known_length([]) == nothing
+
+    x = view(SArray{Tuple{3,3,3}}(ones(3,3,3)), :, SOneTo(2), 2)
+    @test @inferred(ArrayInterface.known_length(x)) == 6
+    @test @inferred(ArrayInterface.known_length(x')) == 6
 end
 
 @testset "indices" begin