Inverting default `collect` behavior when iterator returns an array

A MWE:

struct MyIter
    n
    limit
end

function Base.iterate(m::MyIter)
    state = ones(m.n)
    return  (state,state .+ 1)
end

function Base.iterate(m::MyIter,state)
    if first(state) > m.limit
        return nothing
    else
        new_state = state .+ 1
        return (state,new_state)
    end

end

function Base.length(m::MyIter)
    return m.limit
end


collect(MyIter(2,10))

This will return:

julia> collect(MyIter(2,10))
10-element Vector{Any}:
 [1.0, 1.0]
 [2.0, 2.0]
 [3.0, 3.0]
 [4.0, 4.0]
 [5.0, 5.0]
 [6.0, 6.0]
 [7.0, 7.0]
 [8.0, 8.0]
 [9.0, 9.0]
 [10.0, 10.0]

Is there a different design such that my iterator could return [[1.0,2.0...,10.0],[1.0,2.0...,10.0]] or an n by limitmatrix, without users needing to do that step manually when collecting?

Would something like this work?

struct MyIter
    n
    limit
end

function Base.iterate(m::MyIter)
    state = ones(m.limit)
    for i in 1:m.limit
         state[i] = i
    end
    return  (state, (state,1))
end

function Base.iterate(m::MyIter,state)
    if state[2] > m.n
        return nothing
    else
        return (state[1], state[2]+1)
    end

end

function Base.length(m::MyIter)
    return m.n
end

I’m on mobile on the road so I can’t test it right now, but I think this will return the data in the flipped orientation.

2 Likes

Just realized that this will return arrays that all point to the same object. The second iterate function should return a copy of state[1] if you want all of the numbers to operate independently.

2 Likes

Thank you - I took your essential idea of iterating over each item, therefore effectively collecting column by column instead of row by row.

1 Like