Unroll only the first iteration

[mschauer]

For accumulation of iterators it is nice to unroll only the first iteration. For example

  for (isfirst, x) in flagfirst(xs)
        if isfirst
            y = f(x)
        else
            y = y + f(x)
        end
    end

is a nice idiom and

@unrollfirst for (isfirst, x) in flagfirst(xs)
        if isfirst
            y = f(x)
        else
            y = y + f(x)
        end
    end

would make it type stable.
Would this be easy to add as variation of @unroll, say @unrollfirst? See https://discourse.julialang.org/t/skipping-parts-of-a-for-loop-in-the-first-iteration/16252/19 for related discussion

[cstjean]

I'm not sure I understand what you're asking for. Isn't https://discourse.julialang.org/t/skipping-parts-of-a-for-loop-in-the-first-iteration/16252/8 good enough? It's similar to what I would write for @unrollfirst. I think a more general @unroll_n 5 for ... would be a fine addition to this package.

[cstjean]

I tried my hand at it

     macro unroll_n(n::Int, loop)
           @assert @capture(loop, for var_ in iter0_; body__; end)
           @gensym iter st
           function gen_loop(rem::Int)
               if rem == 0
                   quote
                       while true
                           ϕ === nothing && break
                           $var, $st = ϕ
                           $(body...)
                           ϕ = Base.iterate($iter, $st)
                       end
                   end
               else
                   quote
                       if ϕ !== nothing
                           let ($var, $st) = ϕ
                               $(body...)
                               ϕ = Base.iterate($iter, $st)
                               $(gen_loop(rem-1))
                           end
                       end
                   end
               end
           end
           
           esc(quote
               $iter = $iter0
               ϕ = Base.iterate($iter)
               $(gen_loop(n))
               end)
       end

There's a weird scoping bug I don't understand, unfortunately

julia> @unroll_n 1 for x in 1:10
       println(x)
       end
ERROR: UndefVarError: ϕ not defined
Stacktrace:
 [1] top-level scope at ./REPL[49]:17

But the macroexpansion looks right.

[mschauer]

Could this be related to the observation that https://github.com/mschauer/Trajectories.jl/blob/master/src/unroll1.jl will only run if defined in the same module as the caller?