# Collect an iterator into an SVector or Tuple

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)))
``````

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

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

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

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

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`.

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
``````