Using the iteration interface with custom indices

New to Julia
1

I’m a bit confused on how to use iteration utilities with a custom index set. I tried a MWE (just reproducing a for loop) as follows:

struct CustomIndex
    i::Int
end

struct CustomIndexSet
    n::Int
end

Base.iterate(I::CustomIndexSet) = I,CustomIndex(1)
function Base.iterate(I::CustomIndexSet,i::CustomIndex)
    if i.i+1 <= I.n
        return I,CustomIndex(i.i+1)
    else
        return nothing
    end
end

I want for i in CustomIndexSet(4) to give 1,2,3,4, but instead I get

julia> for i in CustomIndexSet(4)
       @show i
       end
i = CustomIndexSet(4)
i = CustomIndexSet(4)
i = CustomIndexSet(4)
i = CustomIndexSet(4)

What’s the right way to go about this?

2

I also tried just defining a iterator using just CustomIndex (with hardcoded stopping point)

struct CustomIndex
    i::Int
end
Base.iterate(I::CustomIndex) = CustomIndex(1),1
function Base.iterate(i::CustomIndex,state)
    if i.i+1 <= 4
        return CustomIndex(i.i+1),state+1
    else
        return nothing
    end
end

then for i in CustomIndex(1) gives an infinite loop with

i = CustomIndex(1)
i = CustomIndex(1)
...

However, the iterate function seems to work (returns nothing when expected).

3

it’s because the first argument returned by iterate is the “value” (i) you will get:

julia> struct CustomIndex
           n::Int
       end

julia> Base.iterate(S::CustomIndex, state=1) = state > S.n ? nothing : (state, state+1)

julia> for i in CustomIndex(4)
              @show i
              end
i = 1
i = 2
i = 3
i = 4
1 Like
4

This solves the MWE, but I want state to be a CustomIndex. Does this mean that if state is a CustomIndex, then it needs to also carry information about the range?

In my actual problem, I’m trying to increment state as a custom-typed multi-index.

5

in your first example, your I.n is always 4, so no surprise there. And your custom-typed indexing was working just fine.

Base.iterate(I::CustomIndexSet, i::CustomIndex=CustomIndex(1)) = i.i, i

julia> function Base.iterate(I::CustomIndexSet,i::CustomIndex)
        i.i <= I.n && return i.i,CustomIndex(i.i+1)
        nothing
   end

julia> for i in CustomIndexSet(4)
              @show i
           end
i = 1
i = 2
i = 3
i = 4
1 Like
6

I see that I reversed the arguments in

Base.iterate(I::CustomIndexSet) = 1,CustomIndex(1)

I got confused by the description in the manual of iterate:

Returns either a tuple of the first item and initial state or nothing if empty

“Item” and “state” seem interchangeable and ambiguous.

1 Like
7

This doesn’t quite fix it though - I’m trying to extract the state from the iterator, e.g. something like

julia> for i in CustomIndexSet(4)
              @show i
           end
i = CustomIndex(1)
i = CustomIndex(2)
i = CustomIndex(3)
i = CustomIndex(4)
8 
julia> function Base.iterate(I::CustomIndexSet,i::CustomIndex)
           if i.i <= I.n
               return i,CustomIndex(i.i+1)
           else
               return nothing
           end
       end

julia> for i in CustomIndexSet(4)
              @show i
           end
i = CustomIndex(1)
i = CustomIndex(2)
i = CustomIndex(3)
i = CustomIndex(4)

difference on line 3

2 Likes
11

Thanks to @jling, here’s a MWE:

struct CustomIndex
    i::Int
end
struct CustomIndexSet
    n::Int
end
# Base.iterate(I::CustomIndexSet) = 1,CustomIndex(1)
Base.iterate(I::CustomIndexSet, i::CustomIndex=CustomIndex(1)) = i.i, i
function Base.iterate(I::CustomIndexSet,i::CustomIndex)
   if i.i <= I.n
       return i,CustomIndex(i.i+1)
   else
       return nothing
   end
end

which gives

julia> for i in CustomIndexSet(4)
       @show i
       end
i = CustomIndex(1)
i = CustomIndex(2)
i = CustomIndex(3)
i = CustomIndex(4)
1 Like
12

I’m still am unclear on why these two lines return different results:

Base.iterate(I::CustomIndexSet) = 1,CustomIndex(1)
Base.iterate(I::CustomIndexSet, i::CustomIndex=CustomIndex(1)) = i.i, i

The first line forces a return of CustomIndex(1) as behavior for the first iterate

julia> for i in CustomIndexSet(4)
       @show i
       end
i = 1
i = CustomIndex(1)
i = CustomIndex(2)
i = CustomIndex(3)
i = CustomIndex(4)

The second one has CustomIndex(1) as an optional argument - but why does this remove the appearance of an extra entry in the for loop?

13

For anyone reading this thread, @mkitti clarified this on Slack for me.

I was having Base.iterate return current item + next state, when it should just return next item + next state (which can be the same). MWE for this is now:

struct CustomIndex
    i::Int
end
struct CustomIndexSet
    n::Int
end
Base.iterate(I::CustomIndexSet) = CustomIndex(1),CustomIndex(1)
function Base.iterate(I::CustomIndexSet,i::CustomIndex)
    if i.i < I.n
        return CustomIndex(i.i+1),CustomIndex(i.i+1)
    else
       return nothing
   end
end

which gives for

for i in CustomIndexSet(3)
    @show i
end

the desired result

for i in CustomIndexSet(3)
       @show i
       end
i = CustomIndex(1)
i = CustomIndex(2)
i = CustomIndex(3)
1 Like