How to make type stable generator with a while condition like map?

user_371843 0
Viewed 71

Say I'm simulating an experiment and I don't know when the step will terminate. With the while condition and pushing to an array of Any, the function isn't type stable.

function simulate()
    result = []
    alive = true
    while alive
        push!(result,alive)
        alive = rand(Bool)
    end
    return result
end

Is there a way to accomplish this in a type stable way?

julia control-flow list-comprehension types
1 Answers

There is a very easy fix in this particular case! You can specify an empty Array{Bool} instead of the Array{Any} you are currently getting by simply writing Bool[] in place of []:

function simulate()
    result = Bool[]
    alive = true
    while alive
        push!(result,alive)
        alive = rand(Bool)
    end
    return result
end

Looking at the output of a @code_warntype on this modified function, we can see that the result is now stably inferred as a Vector{Bool}

julia> @code_warntype simulate()
MethodInstance for simulate()
  from simulate() in Main at REPL[13]:1
Arguments
  #self#::Core.Const(simulate)
Locals
  alive::Bool
  result::Vector{Bool}
Body::Vector{Bool}
1 ─     (result = Base.getindex(Main.Bool))
│       Main.sizehint!(result, 50)
└──     (alive = true)
2 ┄     goto #4 if not alive
3 ─     Main.push!(result, alive::Core.Const(true))
│       (alive = Main.rand(Main.Bool))
└──     goto #2
4 ─     return result
user_136397340
Related