How to restrict an iteration with StepRange?

Given an iterator (for example from the ‘Interfaces’ doc).

s = Squares(9) 
collect(s) # [1, 4, 9, 16, 25, 36, 49, 64, 81]

Now given a StepRange

R = 1:2:11  # 1, 3, 5, 7, 9, 11 

how can I restrict the iteration over the Squares to the values of R to get the squares restricted to R:

[1, 9, 25, 49, 81]

Neither collect(R, s) nor collect(s, R) works, somewhat surprisingly.

More generally, how can I restrict an iteration to an iteration respecting a given StepRange?

Edit: Steps of the form 0:2:n seems to be a special problem. Both suggestions below lead to errors in this case. Here is a minimal example of what I expect a RangeIterator to do:

struct Squares count::Int end
Base.iterate(S::Squares, state=0) = state > S.count ? nothing : (state*state, state+1)
Base.length(S::Squares) = S.count
EvenBisectSquares(n) = RangeIterator(Squares(n), 0:2:n)
OddBisectSquares(n)  = RangeIterator(Squares(n), 1:2:n)
EvenBisectSquares(9) |> println ∘ collect  # [0, 4, 16, 36, 64]
OddBisectSquares(9)  |> println ∘ collect  # [1, 9, 25, 49, 81]

I don’t think exists but implementing it would be a great exercise to learn how to use the iteration protocol.

Somewhat relatedly, check out RFC: jumping in iterators.

You can always do collect(s)[R]. It’s not optimal, but you don’t always need optimal :slight_smile:

1 Like

Here is a toy implementation (hidden to not spoil the fun):

Summary
struct RestrictIterator{I, R}
    r::R # should iterate ints
    i::I
end

Base.IteratorSize(::Type{<:RestrictIterator}) = Base.SizeUnknown()
Base.IteratorEltype(::Type{<:RestrictIterator{I}}) where{I} = Base.IteratorEltype(I)
Base.eltype(i::RestrictIterator) = eltype(i.i)

macro iterate_helper(iter, sr)
    return quote
        ir = if $(esc(sr)) === nothing
            iterate($(esc(iter)))
        else
            iterate($(esc(iter)), $(esc(sr)))
        end
        ir === nothing && return
        ir[1], ir[2]
    end
end

function Base.iterate(I::RestrictIterator, state=(nothing, nothing, 0))
    si, sr, vr0 = state

    vr1::Integer, sr = @iterate_helper I.r sr
    vr = vr1 - vr0
    @assert vr > 0

    # Replace with "jump"-iteration when that exist
    vi, si = @iterate_helper I.i si
    for i in 1:(vr-1)
        vi, si = @iterate_helper I.i si
    end
    
    return vi, (si, sr, vr1)
end

Usage:

julia> s = Squares(9)
Squares(9)

julia> R = 1:2:11
1:2:11

julia> collect(RestrictIterator(R, s))
5-element Array{Int64,1}:
  1
  9
 25
 49
 81
2 Likes

Great! For me there is only one question left: Why is this a ‘toy implementation’ and not ready to be available as Base.Iterators.RangeIterator in Julia 1.1 from 24/12/2018 onwards? I feel this is highly desirable.

Perhaps a quite uncommon use case? Could be put in https://github.com/JuliaCollections/IterTools.jl?

1 Like

A simple test case is bisection: the odd case 1:2:n and the even case 0:2:n (giving 0, 4, 16,…). The even case runs into the AssertionError: vr > 0.

I left some bugs in there as homework to the reader :wink:

Nice example. I am still solidifying the design, but this can now be expressed by DynamicIterator primitives:

using DynamicIterators

# Assign iterate number as key, like enumerate
Sq = TimeLift(Squares())

# Create NextKeys control slot
C = attach(NextKeys, Sq)

# Bind iterator 4:2:8 to the control slot and collect
collect(bind(4:2:8, C))

This gives

3-element Array{Any,1}:
 4 => 1 
 6 => 9 
 8 => 25

Let’s try to gain some insight and compare two cases:

J(n::Int) = filter(isSquare, 1:3:n)
for i in J(100) print("$i, ") end; println()

R(n::Int) = RestrictIterator(Squares(n), 1:3:n)
for i in R(20) print("$i, ") end; println()

The first case basically iterates over the squares of i if i mod 3 != 0.

for n in 1:20 rem(n, 3) != 0 && n^2 <= 100 && print(n^2, ", ") end

The second case iterates over the squares of i if i mod 3 = 1.

for n in 1:20 rem(n, 3) == 1 && print(n^2, ", ") end

I leave it to you to decide which use case is more uncommon.

But: it is much more expensive to find out if an integer is a square (case 1, ‘isSquare’) than to generate the squares (case 2, ‘Squares’).

(I maintain that this is a general characteristic of an analytical versus a synthetic approach seen in many situations.)

For this reason I think there should be iterators with generating functions and step-ranges in addition to the familiar filter iterator with predicates and step-ranges.