Search code examples
typesjulialist-comprehensioncontrol-flow

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


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?


Solution

  • 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