`HasLength` for even more `Flatten` iterators

[tkoolen]

I saw these recent changes: #22691, but maybe we can do better.

Even after these changes, the following example (from Discourse) still results in a SizeUnknown Flatten iterator:

using Base.Iterators: flatten
dr = [1:2, 4:7]
Base.iteratorsize(flatten(dr)) # SizeUnknown()

It would be possible to simply define

Base.length(fl::Flatten) = mapreduce(length, +, 0, fl.it)

and to add UnitRange to this list, but maybe it would be nice to do this more generically? Base.iteratorsize could return HasLength() for fl::Flatten if all iterators in fl.it are HasLength or HasSize, and SizeUnknown otherwise. That would mean that Base.iteratorsize(::Flatten) could no longer be optimized into a constant however; is that a big issue?

[rfourquet]

Even only with UnitRange, it would mean that computing lenght is not O(1), which I think is a requirement for HasLength. I didn't see it specified though.

[stevengj]

I think it is desirable for iteratorsize to be type stable (i.e. type of result known at compile time given the input type).

I don't think it is necessary for length to be O(1). (It isn't for String, after all.)

[stevengj]

Also, I'm wary of actually using the iterator to compute the length in general, because an arbitrary iterator may have side effects. Think of Base.Iterators.flatten(eachline(filename)), for example.

[tkoolen]

OK, I clearly didn't think it through enough. Would a special case for UnitRange still be desirable, or would you consider it to just be code bloat for little additional gain? I think it may be a common enough case that it could warrant special optimization.

[stevengj]

Couldn't we have a special case for any AbstractArray?

[tkoolen]

I don't think that's possible, because e.g. an OffsetArrays.OffsetArray is a subtype of AbstractArray, but does not define length or size methods.

This whole thing makes me wish that the interface were to define methods for Base.iteratorsize(::Type{T}) instead of for Base.iteratorsize(T); I believe everything could be done in type stable fashion using Tuple type recursion if that were the case.

[mcabbott]

Defining length for more Flatten types could greatly speed up collect ∘ Iterators.flatten. This has been discussed a bit in the context of #43334, and how best to express things like reduce(vcat, vec_of_vecs). Times:

julia> vecs = [rand(100) for _ in 1:100];

julia> @btime collect(Iterators.flatten($vecs));
  min 68.875 μs, mean 81.634 μs (9 allocations, 326.55 KiB) # master
  min 12.167 μs, mean 17.848 μs (2 allocations, 78.17 KiB)  # with length(::Flatten{...

I presume much has changed since discussion above in 2017, and in particular OffsetArrays do define length now, and Base.IteratorSize does work on types.

[zot]

I made Append as a basis for intersperse (which I first wrote in terms of generators but then ran into this issue).

Maybe this code is naive but perhaps it could be a fallback for when Flatten can't determine length?

Anyway, it's a strawman, at least.

Here's the code

intersperse(iterator, sep) =
    append((
        (first(iterator),),
        append(Iterators.zip(
            Iterators.repeated(sep, length(iterator) - 1),
            Iterators.drop(iterator, 1)))
    ))

# Append several iterators together

struct Append{Length,ElType}
    it
end

"""
    append(iters)

Given an iterator that yields iterators, return an iterator that yields the
elements of those iterators.
"""
function append(iters)
    tot = 0
    typ = Any
    for i in iters
        if tot >= 0
            if Base.IteratorSize(typeof(i)) == Base.HasLength()
                tot += length(i)
            else
                tot = -1
            end
        end
        if !isnothing(typ) && Base.IteratorEltype(typeof(i)) == Base.HasEltype()
            typ = promote_type(typ, eltype(i))
        else
            typ = nothing
        end
    end
    return Append{tot >= 0 ? tot : :Unknown,typ}(iters)
end

Base.IteratorSize(j::Type{Append}) = Base.HasLength()
Base.IteratorSize(j::Type{Append{:Unknown}}) = Base.SizeUnknown()

Base.IteratorEltype(j::Type{Append{T,nothing}}) where {T} = Base.EltypeUnknown()
Base.IteratorEltype(j::Type{Append{T}}) where {T} = Base.HasEltype

Base.length(::Append{N}) where {N} = N isa Integer ? N : throw(ArgumentError("Unknown length"))

Base.eltype(j::Append{T,U}) where {T,U} = U
Base.eltype(j::Append{T,nothing}) where {T} = throw(ArgumentError("Unknown eltype"))

Base.iterate(j::Append) = isempty(j.it) ? nothing :
                          Base.iterate(j, (first(j.it), Iterators.drop(j.it, 1)))
function Base.iterate(j::Append, (f, r))
    isempty(f) && isempty(r) &&
        return nothing
    isempty(f) &&
        return Base.iterate(j, (first(r), Iterators.drop(r, 1)))
    return first(f), (Iterators.drop(f, 1), r)
end