Iterate on calculated iterable

General Usage
,
1

I have a type for which I want to implement a kind of iteration, which I can easily do with functions in Iterators. But I want iteration to happen on something I construct. MWE:

struct NamedProduct{T}
    np::T
end

function make_iterable(itr::NamedProduct{<:NamedTuple{N}}) where N
    Iterators.map(NamedTuple{N}, Iterators.product(values(itr.np)...))
end

I = NamedProduct((a = 1:4, b = 2:3))
collect(make_iterable(I))
collect(I)                      # <= I want this to be equivalent

So one horrible, horrible solution I came up with is hiding the resulting iterator in the state, like this:

Base.IteratorSize(itr::NamedProduct) = Base.HasShape{length(itr.np)}()

# assuming inner iterators are flat, for simpler code
Base.axes(itr::NamedProduct) = map(itr -> axes(itr, 1), values(itr.np))

function Base.iterate(itr::NamedProduct)
    inner_itr = make_iterable(itr)
    r = iterate(inner_itr)
    r ≡ nothing && return r
    elt, inner_state = r
    elt, (inner_itr, inner_state)
end

# these two functions could be combined into one,
# here they are two for clarity
function Base.iterate(itr::NamedProduct, (inner_itr, inner_state))
    r = iterate(inner_itr, inner_state)
    r ≡ nothing && return r
    elt, inner_state = r
    elt, (inner_itr, inner_state)
end

I am wondering if there is a package that implements a wrapper that does this, so I would not reinvent the wheel, just use that and keep my code simpler.

Also, if there is a better (nicer? more elegant?) way of doing it.

1 Like
2

Although you probably already thought of this, trivially you could do

Base.iterate(n::NamedProduct, s...) = iterate(make_iterable(n), s...)

However, this reconstructs the wrapped iterator for every iteration, which might or might not be acceptable.

Also you could add an additional field to NamedProduct that holds the iterator and handle construction transparently in the constructor.

I would also consider a small variation of your approach that partly avoids explicit state management (and is a little less horrible).

forward_itr(_, ::Nothing, ) = nothing
forward_itr(inner_itr, (el, state)) = el, (inner_itr, state)

function Base.iterate(itr::NamedProduct)
    inner_itr = make_iterable(itr)
    return forward_itr(inner_itr, iterate(inner_itr))
end

function Base.iterate(itr::NamedProduct, (inner_itr, inner_state))
    return forward_itr(inner_itr, iterate(inner_itr, inner_state))
end

Unfortunately, I can’t think of a package that does something like this.

2 Likes
3

Thanks, I find it rather neat and it makes the code much more readable!

I am thinking about a PR to IterTools.jl, but I am wondering how this could be packaged best. Probably a macro with the syntax

@forward_iterator itr::NamedProduct make_iterable(itr)

which then uses the (not necessarily exposed) forward_itr methods.

1 Like
4

This may be necessary, if one wants to handle the various iterator traits by forwarding transparently.

5

Something like works fairly well:

struct NamedProduct{T,I}
    np::T
    itr::I
    function NamedProduct(np::NamedTuple{N}) where N
        itr = Iterators.map(NamedTuple{N}, Iterators.product(values(np)...))
        new{typeof(np), typeof(itr)}(np, itr)
    end
end

Base.show(io::IO, I::NamedProduct) = show(io, I.np)

Base.IteratorSize(itr::NamedProduct) = Base.HasShape{length(itr.np)}()
Base.axes(itr::NamedProduct) = map(itr -> axes(itr, 1), values(itr.np))

Base.iterate(itr::NamedProduct, state...) = iterate(itr.itr, state...)

I = NamedProduct((a = 1:4, b = 2:3))
collect(I)

If making a general thing with a macro as indicated, it’s perhaps better to make a named tuple containing the original NamedProduct struct and the itr as separate fields, with a unique field name, so you can dispatch on it. I suppose you need some extra information for the IteratorSize and axes for a completely general macro.

1 Like
6

Yes, this precisely captures my suggestion 2 from above.
It should also work to just forward the calls to Base.IteratorSize and Base.axes to the itr field.
This would also make implementing a general macro far easier - simply require the created iterator type to implement the iterator interface.

This approach certainly shines due to its simplicity.
However, it also requires to change the struct layout and introduces an additional type parameter.

What makes my final approach above in my opinion so appealing, is that there is no need to modify NamedProduct.
It basically provides a template to plug in an arbitrarily transformed version of an object for iteration, without touching the rest of the code.

2 Likes
7

I agree that modifying NamedProduct can be messy. It’s not too difficult to use MacroTools.jl to parse the struct definition to add another parameter and a field, but it can get messy with constructors and traits.

I.e. something along the suggestion above, but at the struct definition level:

@forward_iterator struct NamedProduct{T} np::T end makeiter

should ideally rewrite it to

struct NamedProduct{T, I}
    np::T
    itr::I
    function NamedProduct(np::T) where T  # this signature depends on ... uh ... the original struct
        it = makeiter(np)
        new{T, typeof(it)}(np, it)
    end
end
Base.iterate(np::NamedProduct, state...) = iterate(np.itr, state...)

The modified struct could be made a subtype of some AbstractForwardIter which has the iterate function.