coax inference

Author: chriselrod

src/ArrayInterface.jl

@@ -94,14 +94,27 @@ known_length(::Type{<:Number}) = 1
 known_length(::Type{<:AbstractCartesianIndex{N}}) where {N} = N
 known_length(::Type{T}) where {T} = _maybe_known_length(Base.IteratorSize(T), T)
 
-@inline _prod_or_nothing(x, ::Tuple{}) = x
-@inline _prod_or_nothing(_, ::Tuple{Nothing,Vararg}) = nothing
-@inline _prod_or_nothing(x, y::Tuple{I,Vararg}) where {I} = _prod_or_nothing(x*first(y), Base.tail(y))
+@generated function _prod_or_nothing(x::Tuple)
+  p = 1
+  for i in eachindex(x.parameters)
+    x.parameters[i] === Nothing && return nothing
+    p *= x.parameters[i].parameters[1]
+  end
+  StaticInt(p)
+end
+
+function _maybe_known_length(::Base.HasShape, ::Type{T}) where {T}
+  t = map(_static_or_nothing, known_size(T))
+  _int_or_nothing(_prod_or_nothing(t))
+end
 
-_maybe_known_length(::Base.HasShape, ::Type{T}) where {T} = _prod_or_nothing(1, known_size(T))
 _maybe_known_length(::Base.IteratorSize, ::Type) = nothing
+_static_or_nothing(::Nothing) = nothing
+@inline _static_or_nothing(x::Int) = StaticInt{x}()
+_int_or_nothing(::StaticInt{N}) where {N} = N
+_int_or_nothing(::Nothing) = nothing
 function known_length(::Type{<:Iterators.Flatten{I}}) where {I}
-  _prod_or_nothing(1, (known_length(I),known_length(eltype(I))))
+  _int_or_nothing(_prod_or_nothing((_static_or_nothing(known_length(I)),_static_or_nothing(known_length(eltype(I))))))
 end
 
 """

test/runtests.jl

@@ -518,11 +518,7 @@ end
 
     @test @inferred(ArrayInterface.size(irev)) === (StaticInt(2), StaticInt(3), StaticInt(4))
     @test @inferred(ArrayInterface.size(iprod)) === (StaticInt(2), StaticInt(3), StaticInt(4))
-    if VERSION >= v"1.7"
-      @test @inferred(ArrayInterface.size(iflat)) === (static(72),)
-    else
-      @test_skip @inferred(ArrayInterface.size(iflat)) === (static(72),)
-    end
+    @test @inferred(ArrayInterface.size(iflat)) === (static(72),)
     @test @inferred(ArrayInterface.size(igen)) === (StaticInt(2), StaticInt(3), StaticInt(4))
     @test @inferred(ArrayInterface.size(iacc)) === (StaticInt(2), StaticInt(3), StaticInt(4))
     @test @inferred(ArrayInterface.size(ienum)) === (StaticInt(2), StaticInt(3), StaticInt(4))