# Using the iteration interface with custom indices

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

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

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

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.

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

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

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

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

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?

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