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

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

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)

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)

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)