Collect an iterator into an SVector or Tuple

General Usage
1

Is there a recommended idiom for collecting a number of elements (known at compile time, embedded in some type parameter) from an iterable into a Tuple or SVector?

I can write a generated function and unroll, but I was hoping for something simple and clean. MWE:

using StaticArrays
struct Fibonacci end
function Base.iterate(::Fibonacci, state = (0, 1))
    f1, f2 = state
    f0 = f1 + f2
    f0, (f0, f1)
end
Base.IteratorEltype(::Type{Fibonacci}) = Base.HasEltype()
Base.eltype(::Type{Fibonacci}) = Int
Base.IteratorSize(::Type{Fibonacci}) = Base.IsInfinite()

# want the result here, without collecting first
SVector{5}(collect(Iterators.take(Fibonacci(), 5)))
2

Here’s my attempt. Given an iterator x, we can generate a tuple by repeatedly calling iterate and taking the values:

julia> x = [1, 2, 3];

julia> (((val, state) = iterate(x); val), ((val, state) = iterate(x, state); val), ((val, state) = iterate(x, state); val))
(1, 2, 3)

Here’s a macro that automates that process:

julia> macro tuple_unpack(N::Int, x)
         @assert N >= 1
         expr = Expr(:tuple)
         push!(expr.args, quote
           begin
             (val, state) = iterate($(esc(x)))
             val
           end
         end)
         for i = 2:N
           push!(expr.args, quote
             begin
               (val, state) = iterate($(esc(x)), state)
               val
             end
           end)
         end
         expr
       end
@tuple_unpack (macro with 1 method)

Usage:

julia> x = [1, 2, 3];

julia> @tuple_unpack 3 x
(1, 2, 3)

julia> g = (2 * i for i in 1:4)
Base.Generator{UnitRange{Int64},var"#13#14"}(var"#13#14"(), 1:4)

julia> @tuple_unpack 4 g
(2, 4, 6, 8)

Performance seems good, at least for a very simple case:

julia> f(x) = @tuple_unpack 3 x
f (generic function with 1 method)

julia> @btime f($x)
  6.606 ns (0 allocations: 0 bytes)
(1, 2, 3)
1 Like
3

Thanks, I am looking for something that does not involve metaprogramming.

Behold

function static_take(::Val{K}, itr) where K
    K == 0 && error("so how would I know the type?")
    dummy = SVector(ntuple(_ -> nothing, Val(K + 1)))
    function f(::Nothing, ::Nothing)
        x, state = iterate(itr)
        SVector(x), state
    end
    function f((xs, state), ::Nothing)
        x, state = iterate(itr, state)
        push(xs, x), state
    end
    xs, _ = foldl(f, dummy)
    SVector(xs)
end

static_take(Val(5), Fibonacci())

which is not going to win any beauty contests, but it infers (but of course in a way this is cheating, since StaticArrays does the code generation).

2 Likes
4

You can do

using BangBang: push!!
foldl(push!!, itr; init=SA[])

which can be generalized to work with arbitrary (well-behaving) array type:

using BangBang: push!!, Empty
collectto(T::Type, itr) = foldl(push!!, itr; init=Empty(T))
1 Like
5

Actually, if you just need it for SVector, maybe the easiest thing to do is:

mapfoldl(x -> SA[x], vcat, itr; init=SA[])
2 Likes
6

This is indeed easy, but it does not infer the size. Eg

take2(::Val{K}, itr) where K = mapfoldl(x -> SA[x], vcat, Iterators.take(itr, K); init=SA[])
take2(Val{5}(), Fibonacci())

infers as Any.

7

I just realized that I can use a Tuple for counting (and nothing else), resulting in this nice example of compiler abuse.

_rec(acc, itr, state, token1) = acc
function _rec(acc, itr, state, token1, tokens...)
    x, state′ = iterate(itr, state)
    _rec((acc..., x), itr, state′, tokens...)
end
function take3(cnt, itr)
    r = iterate(itr)
    r ≡ nothing && return ()
    _rec((first(r), ), itr, last(r), ntuple(_ -> nothing, cnt)...)
end

@code_warntype take3(Val{5}(), Fibonacci()) # infers