Support static info on iterators (#236)

* Support static info on iterators

* New support for types from `Base.Iterators`
* Check for `Base.IteratorSize` before assuming we can rely on axes when
  finding the size.
* Update internal calculation of range length to match base

* Use `Iterators.Pairs` for v1.6 support

* Test size error and Iterators.product

* Use fall backs to `static_length` and `known_length`

Use length methods for all non `HasShape` iterators and error or
evaluate there

Author: Tokazama

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "4"
+version = "4.0.1"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/ArrayInterface.jl

@@ -90,15 +90,13 @@ known_length(x) = known_length(typeof(x))
 known_length(::Type{<:NamedTuple{L}}) where {L} = length(L)
 known_length(::Type{T}) where {T<:Slice} = known_length(parent_type(T))
 known_length(::Type{<:Tuple{Vararg{Any,N}}}) where {N} = N
-known_length(::Type{T}) where {Itr,T<:Base.Generator{Itr}} = known_length(Itr)
 known_length(::Type{<:Number}) = 1
 known_length(::Type{<:AbstractCartesianIndex{N}}) where {N} = N
-function known_length(::Type{T}) where {T}
-    if parent_type(T) <: T
-        return missing
-    else
-        return prod(known_size(T))
-    end
+known_length(::Type{T}) where {T} = _maybe_known_length(Base.IteratorSize(T), T)
+_maybe_known_length(::Base.HasShape, ::Type{T}) where {T} = prod(known_size(T))
+_maybe_known_length(::Base.IteratorSize, ::Type) = missing
+function known_length(::Type{<:Iterators.Flatten{I}}) where {I}
+    known_length(I) * known_length(eltype(I))
 end
 
 """

src/dimensions.jl

@@ -3,6 +3,20 @@ function throw_dim_error(@nospecialize(x), @nospecialize(dim))
     throw(DimensionMismatch("$x does not have dimension corresponding to $dim"))
 end
 
+@propagate_inbounds function _promote_shape(a::Tuple{A,Vararg{Any}}, b::Tuple{B,Vararg{Any}}) where {A,B}
+    (_try_static(getfield(a, 1), getfield(b, 1)), _promote_shape(tail(a), tail(b))...)
+end
+_promote_shape(::Tuple{}, ::Tuple{}) = ()
+@propagate_inbounds function _promote_shape(::Tuple{}, b::Tuple{B}) where {B}
+    (_try_static(static(1), getfield(b, 1)),)
+end
+@propagate_inbounds function _promote_shape(a::Tuple{A}, ::Tuple{}) where {A}
+    (_try_static(static(1), getfield(a, 1)),)
+end
+@propagate_inbounds function Base.promote_shape(a::Tuple{Vararg{CanonicalInt}}, b::Tuple{Vararg{CanonicalInt}})
+    _promote_shape(a, b)
+end
+
 #julia> @btime ArrayInterface.is_increasing(ArrayInterface.nstatic(Val(10)))
 #  0.045 ns (0 allocations: 0 bytes)
 #ArrayInterface.True()

src/ranges.jl

@@ -304,66 +304,48 @@ end
     val::Int
 end
 
+@noinline unequal_error(x,y) = @assert false "Unequal Indices: x == $x != $y == y"
+@inline check_equal(x, y) = x == y || unequal_error(x,y)
+_try_static(::Missing, ::Missing) = missing
+_try_static(x::Int, ::Missing) = x
+_try_static(::Missing, y::Int) = y
 @inline _try_static(::StaticInt{N}, ::StaticInt{N}) where {N} = StaticInt{N}()
 @inline function _try_static(::StaticInt{M}, ::StaticInt{N}) where {M,N}
     @assert false "Unequal Indices: StaticInt{$M}() != StaticInt{$N}()"
 end
-@noinline unequal_error(x,y) = @assert false "Unequal Indices: x == $x != $y == y"
-@inline function check_equal(x, y)
-    x == y || unequal_error(x,y)
-end
-@propagate_inbounds function _try_static(::StaticInt{N}, x) where {N}
-    @boundscheck check_equal(StaticInt{N}(), x)
-    return StaticInt{N}()
-end
-@propagate_inbounds function _try_static(x, ::StaticInt{N}) where {N}
-    @boundscheck check_equal(x, StaticInt{N}())
-    return StaticInt{N}()
-end
+@propagate_inbounds _try_static(::StaticInt{N}, x) where {N} = static(_try_static(N, x))
+@propagate_inbounds _try_static(x, ::StaticInt{N}) where {N} = static(_try_static(N, x))
 @propagate_inbounds function _try_static(x, y)
     @boundscheck check_equal(x, y)
     return x
 end
 
 ## length
-@inline function known_length(::Type{T}) where {T<:OptionallyStaticUnitRange}
-    return _range_length(known_first(T), known_last(T))
-end
-
-@inline function known_length(::Type{T}) where {T<:OptionallyStaticStepRange}
-    _range_length(known_first(T), known_step(T), known_last(T))
-end
-
 Base.lastindex(x::OptionallyStaticRange) = length(x)
-Base.length(r::OptionallyStaticUnitRange) = _range_length(static_first(r), static_last(r))
-@inline function Base.length(r::OptionallyStaticStepRange)
+@inline function Base.length(r::OptionallyStaticUnitRange)
     if isempty(r)
         return 0
     else
-        return _range_length(static_first(r), static_step(r), static_last(r))
+        return last(r) - first(r) + 1
     end
 end
-_range_length(start, stop) = missing
-function _range_length(start::CanonicalInt, stop::CanonicalInt)
-    if start > stop
-        return 0
-    else
-        return Int((stop - start) + 1)
-    end
-end
-_range_length(start::CanonicalInt, ::One, stop::CanonicalInt) = _range_length(start, stop)
-@inline function _range_length(start::CanonicalInt, step::CanonicalInt, stop::CanonicalInt)
-    if step > 1
-        return Base.checked_add(Int(div(unsigned(stop - start), step)), 1)
-    elseif step < -1
-        return Base.checked_add(Int(div(unsigned(start - stop), -step)), 1)
-    elseif step > 0
-        return Base.checked_add(Int(div(Base.checked_sub(stop, start), step)), 1)
+Base.length(r::OptionallyStaticStepRange) = _range_length(first(r), step(r), last(r))
+_range_length(start, s, stop) = missing
+@inline function _range_length(start::Int, s::Int, stop::Int)
+   if s > 0
+        if stop < start  # isempty
+            return 0
+        else
+            return Int(div(stop - start, s)) + 1
+        end
     else
-        return Base.checked_add(Int(div(Base.checked_sub(start, stop), -step)), 1)
+        if stop > start  # isempty
+            return 0
+        else
+            return Int(div(start - stop, -s)) + 1
+        end
     end
 end
-_range_length(start, step, stop) = missing
 
 Base.AbstractUnitRange{Int}(r::OptionallyStaticUnitRange) = r
 function Base.AbstractUnitRange{T}(r::OptionallyStaticUnitRange) where {T}

src/size.jl

@@ -16,7 +16,9 @@ julia> ArrayInterface.size(A)
 (static(3), static(4))
 ```
 """
-@inline size(A) = map(static_length, axes(A))
+size(a::A) where {A} = _maybe_size(Base.IteratorSize(A), a)
+_maybe_size(::Base.HasShape{N}, a::A) where {N,A} = map(static_length, axes(a))
+_maybe_size(::Base.HasLength, a::A) where {A} = (static_length(a),)
 size(x::SubArray) = eachop(_sub_size, to_parent_dims(x), x.indices)
 _sub_size(x::Tuple, ::StaticInt{dim}) where {dim} = static_length(getfield(x, dim))
 @inline size(B::VecAdjTrans) = (One(), length(parent(B)))
@@ -40,6 +42,14 @@ function size(a::ReinterpretArray{T,N,S,A}) where {T,N,S,A}
 end
 size(A::ReshapedArray) = Base.size(A)
 size(A::AbstractRange) = (static_length(A),)
+size(x::Base.Generator) = size(getfield(x, :iter))
+size(x::Iterators.Reverse) = size(getfield(x, :itr))
+size(x::Iterators.Enumerate) = size(getfield(x, :itr))
+size(x::Iterators.Accumulate) = size(getfield(x, :itr))
+size(x::Iterators.Pairs) = size(getfield(x, :itr))
+@inline function size(x::Iterators.ProductIterator)
+    eachop(_sub_size, nstatic(Val(ndims(x))), getfield(x, :iterators))
+end
 
 size(a, dim) = size(a, to_dims(a, dim))
 size(a::Array, dim::Integer) = Base.arraysize(a, convert(Int, dim))
@@ -63,6 +73,7 @@ function size(A::SubArray, dim::Integer)
         return static_length(A.indices[pdim])
     end
 end
+size(x::Iterators.Zip) = Static.reduce_tup(promote_shape, map(size, getfield(x, :is)))
 
 """
     known_size(::Type{T}) -> Tuple
@@ -73,7 +84,34 @@ compile time. If a dimension does not have a known size along a dimension then `
 returned in its position.
 """
 known_size(x) = known_size(typeof(x))
-known_size(::Type{T}) where {T} = eachop(_known_size, nstatic(Val(ndims(T))), axes_types(T))
+function known_size(::Type{T}) where {T<:AbstractRange}
+    (_range_length(known_first(T), known_step(T), known_last(T)),)
+end
+known_size(::Type{<:Base.Generator{I}}) where {I} = known_size(I)
+known_size(::Type{<:Iterators.Reverse{I}}) where {I} = known_size(I)
+known_size(::Type{<:Iterators.Enumerate{I}}) where {I} = known_size(I)
+known_size(::Type{<:Iterators.Accumulate{<:Any,I}}) where {I} = known_size(I)
+known_size(::Type{<:Iterators.Pairs{<:Any,<:Any,I}}) where {I} = known_size(I)
+@inline function known_size(::Type{<:Iterators.ProductIterator{T}}) where {T}
+    eachop(_known_size, nstatic(Val(known_length(T))), T)
+end
+
+# 1. `Zip` doesn't check that its collections are compatible (same size) at construction,
+#   but we assume as much b/c otherwise it will error while iterating. So we promote to the
+#   known size if matching a `Missing` and `Int` size.
+# 2. `promote_shape(::Tuple{Vararg{CanonicalInt}}, ::Tuple{Vararg{CanonicalInt}})` promotes
+#   trailing dimensions (which must be of size 1), to `static(1)`. We want to stick to
+#   `Missing` and `Int` types, so we do one last pass to ensure everything is dynamic
+@inline function known_size(::Type{<:Iterators.Zip{T}}) where {T}
+    dynamic(reduce_tup(_promote_shape, eachop(_unzip_size, nstatic(Val(known_length(T))), T)))
+end
+_unzip_size(::Type{T}, n::StaticInt{N}) where {T,N} = known_size(field_type(T, n))
+
+known_size(::Type{T}) where {T} = _maybe_known_size(Base.IteratorSize(T), T)
+function _maybe_known_size(::Base.HasShape{N}, ::Type{T}) where {N,T}
+    eachop(_known_size, nstatic(Val(N)), axes_types(T))
+end
+_maybe_known_size(::Base.IteratorSize, ::Type{T}) where {T} = (known_length(T),)
 _known_size(::Type{T}, dim::StaticInt) where {T} = known_length(field_type(T, dim))
 @inline known_size(x, dim) = known_size(typeof(x), dim)
 @inline known_size(::Type{T}, dim) where {T} = known_size(T, to_dims(T, dim))
@@ -84,3 +122,4 @@ _known_size(::Type{T}, dim::StaticInt) where {T} = known_length(field_type(T, di
         return known_size(T)[dim]
     end
 end
+

test/runtests.jl

@@ -479,8 +479,11 @@ end
     A = zeros(3, 4, 5);
     A[:] = 1:60
     Ap = @view(PermutedDimsArray(A,(3,1,2))[:,1:2,1])';
-    S = @SArray zeros(2,3,4); Sp = @view(PermutedDimsArray(S,(3,1,2))[2:3,1:2,:]);
-    M = @MArray zeros(2,3,4); Mp = @view(PermutedDimsArray(M,(3,1,2))[:,2,:])';
+    S = @SArray zeros(2,3,4)
+    A_trailingdim = zeros(2,3,4,1)
+    Sp = @view(PermutedDimsArray(S,(3,1,2))[2:3,1:2,:]);
+    M = @MArray zeros(2,3,4)
+    Mp = @view(PermutedDimsArray(M,(3,1,2))[:,2,:])';
     Sp2 = @view(PermutedDimsArray(S,(3,2,1))[2:3,:,:]);
     Mp2 = @view(PermutedDimsArray(M,(3,1,2))[2:3,:,2])';
     D = @view(A[:,2:2:4,:]);
@@ -490,6 +493,16 @@ end
     A2 = zeros(4, 3, 5)
     A2r = reinterpret(ComplexF64, A2)
 
+    irev = Iterators.reverse(S)
+    igen = Iterators.map(identity, S)
+    iacc = Iterators.accumulate(+, S)
+    iprod = Iterators.product(axes(S)...)
+    iflat = Iterators.flatten(iprod)
+    ienum = enumerate(S)
+    ipairs = pairs(S)
+    izip = zip(S,S)
+
+
     sv5 = @SVector(zeros(5)); v5 = Vector{Float64}(undef, 5);
     @test @inferred(ArrayInterface.size(sv5)) === (StaticInt(5),)
     @test @inferred(ArrayInterface.size(v5)) === (5,)
@@ -503,6 +516,16 @@ end
     @test @inferred(ArrayInterface.size(A2)) === (4,3,5)
     @test @inferred(ArrayInterface.size(A2r)) === (2,3,5)
 
+    @test @inferred(ArrayInterface.size(irev)) === (StaticInt(2), StaticInt(3), StaticInt(4))
+    @test @inferred(ArrayInterface.size(iprod)) === (StaticInt(2), StaticInt(3), StaticInt(4))
+    @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))
+    @test @inferred(ArrayInterface.size(ipairs)) === (StaticInt(2), StaticInt(3), StaticInt(4))
+    @test @inferred(ArrayInterface.size(izip)) === (StaticInt(2), StaticInt(3), StaticInt(4))
+    @test @inferred(ArrayInterface.size(zip(S, A_trailingdim))) === (StaticInt(2), StaticInt(3), StaticInt(4), static(1))
+    @test @inferred(ArrayInterface.size(zip(A_trailingdim, S))) === (StaticInt(2), StaticInt(3), StaticInt(4), static(1))
     @test @inferred(ArrayInterface.size(S)) === (StaticInt(2), StaticInt(3), StaticInt(4))
     @test @inferred(ArrayInterface.size(Sp)) === (2, 2, StaticInt(3))
     @test @inferred(ArrayInterface.size(Sp2)) === (2, StaticInt(3), StaticInt(2))
@@ -536,6 +559,18 @@ end
     @test @inferred(ArrayInterface.known_size(A2)) === (missing, missing, missing)
     @test @inferred(ArrayInterface.known_size(A2r)) === (missing, missing, missing)
 
+    @test @inferred(ArrayInterface.known_size(irev)) === (2, 3, 4)
+    @test @inferred(ArrayInterface.known_size(igen)) === (2, 3, 4)
+    @test @inferred(ArrayInterface.known_size(iprod)) === (2, 3, 4)
+    @test @inferred(ArrayInterface.known_size(iflat)) === (72,)
+    @test @inferred(ArrayInterface.known_size(iacc)) === (2, 3, 4)
+    @test @inferred(ArrayInterface.known_size(ienum)) === (2, 3, 4)
+    @test @inferred(ArrayInterface.known_size(izip)) === (2, 3, 4)
+    @test @inferred(ArrayInterface.known_size(ipairs)) === (2, 3, 4)
+    @test @inferred(ArrayInterface.known_size(zip(S, A_trailingdim))) === (2, 3, 4, 1)
+    @test @inferred(ArrayInterface.known_size(zip(A_trailingdim, S))) === (2, 3, 4, 1)
+    @test @inferred(ArrayInterface.known_length(Iterators.flatten(((x,y) for x in 0:1 for y in 'a':'c')))) === missing
+
     @test @inferred(ArrayInterface.known_size(S)) === (2, 3, 4)
     @test @inferred(ArrayInterface.known_size(Wrapper(S))) === (2, 3, 4)
     @test @inferred(ArrayInterface.known_size(Sp)) === (missing, missing, 3)
@@ -544,7 +579,6 @@ end
     @test @inferred(ArrayInterface.known_size(Sp2, StaticInt(1))) === missing
     @test @inferred(ArrayInterface.known_size(Sp2, StaticInt(2))) === 3
     @test @inferred(ArrayInterface.known_size(Sp2, StaticInt(3))) === 2
-
     @test @inferred(ArrayInterface.known_size(M)) === (2, 3, 4)
     @test @inferred(ArrayInterface.known_size(Mp)) === (3, 4)
     @test @inferred(ArrayInterface.known_size(Mp2)) === (2, missing)