How to get rid of allocations in this simple filtered generator example?

Here’s a MWE of the allocating code (Julia 1.11.2):

function foo()
    arr = [Ref(i) for i=1:100]
    iter = (i for i in arr if rand()>0)
    @time begin
        elem, state = iterate(iter)
        while true
            next = iterate(iter, state)
            isnothing(next) && break
            elem, state = next
        end
    end
    return nothing
end
foo()

which will show:

  0.000003 seconds (200 allocations: 6.250 KiB)

If you make the element bitstype/immutable (Ref(i)i) it becomes non-allocating, or if you remove the if rand()>0 it also becomes non-allocating. But the combination (which is unavoidable in my real code) apparently allocates every single element twice.

Any suggestions how to fix this?

The allocations go away if you don’t use the global rng state:

using Random

function foo()
    rng = Xoshiro(0)
    arr = [Ref(i) for i=1:100]
    iter = (i for i in arr if rand(rng)>0)
    @time begin
        elem, state = iterate(iter)
        while true
            next = iterate(iter, state)
            isnothing(next) && break
            elem, state = next
        end
    end
    return nothing
end
julia> foo()
  0.000000 seconds
1 Like

Cool.

  1. Is there any special method which can point this out quickly?
  2. The rng allocations are because of per thread state? Asking because, normally one wouldn’t expect some global state to allocate.