How to write a GroupBy iterator: bad and good (requires IterTools at branch ed/iterate)

bad.jl

using IterTools
using IterTools: PeekIter
import Base: iterate, eltype
using Base: IteratorSize, SizeUnknown

macro something(ex)
    quote
        result = $(esc(ex))
        result === nothing && return nothing
        result
    end
end

struct GroupBy2{I, F<:Base.Callable}
    xs::PeekIter{Base.Generator{I, F}}
end

IteratorSize(::Type{<:GroupBy2}) = SizeUnknown()

eltype(::Type{<:GroupBy2{I}}) where {I} = Vector{eltype(I)}

function groupby2(keyfunc::F, xs::I) where {F<:Base.Callable, I}
    keyval(x) = (keyfunc(x), x)

    gen = Base.Generator(keyval, xs)

    GroupBy2{I, typeof(keyval)}(PeekIter(gen))
end

function iterate(it::GroupBy2{I, F}, state=()) where {I, F<:Base.Callable}
    values = Vector{eltype(I)}()

    ((key, val), xs_state) = @something iterate(it.xs, state...)
    push!(values, val)
    peeked = peek(it.xs, xs_state)
    while peeked !== nothing && coalesce(peeked)[1] == key
        ((_, val), xs_state) = iterate(it.xs, xs_state)
        push!(values, val)
        peeked = peek(it.xs, xs_state)
    end

    return (values, (xs_state,))
end

Benchmarking

julia> @benchmark collect(groupby2(first, ["face", "foo", "bar", "book", "baz"]))
BenchmarkTools.Trial:
  memory estimate:  1.98 KiB
  allocs estimate:  60
  --------------
  minimum time:     14.258 μs (0.00% GC)
  median time:      15.473 μs (0.00% GC)
  mean time:        16.915 μs (0.00% GC)
  maximum time:     207.366 μs (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     1
  
julia> @benchmark collect(groupby1(first, ["face", "foo", "bar", "book", "baz"]))
BenchmarkTools.Trial:
  memory estimate:  656 bytes
  allocs estimate:  12
  --------------
  minimum time:     299.350 ns (0.00% GC)
  median time:      314.264 ns (0.00% GC)
  mean time:        412.854 ns (20.56% GC)
  maximum time:     202.104 μs (99.77% GC)
  --------------
  samples:          10000
  evals/sample:     254

Definition.md

groupby(f, xs)

Group consecutive values that share the same result of applying f.

julia> for i in groupby(x -> x[1], ["face", "foo", "bar", "book", "baz", "zzz"])
           @show i
       end
i = String["face", "foo"]
i = String["bar", "book", "baz"]
i = String["zzz"]

good.jl

using IterTools
import Base: iterate, eltype
using Base: IteratorSize, SizeUnknown

macro something(ex)
    quote
        result = $(esc(ex))
        result === nothing && return nothing
        result
    end
end

macro something(ex)
    quote
        result = $(esc(ex))
        result === nothing && return nothing
        result
    end
end

struct GroupBy1{I, F<:Base.Callable}
    xs::I
    keyfunc::F
end
IteratorSize(::Type{<:GroupBy1}) = SizeUnknown()

eltype(::Type{<:GroupBy1{I}}) where {I} = Vector{eltype(I)}

function groupby1(keyfunc::F, xs::I) where {F<:Base.Callable, I}
    GroupBy1{I, F}(xs, keyfunc)
end

function iterate(it::GroupBy1{I, F}, state=nothing) where {I, F<:Base.Callable}
    if state === nothing
        prev_val, xs_state = @something iterate(it.xs)
        prev_key = it.keyfunc(prev_val)
        keep_going = true
    else
        keep_going, prev_key, prev_val, xs_state = state
        keep_going || return nothing
    end
    values = Vector{eltype(I)}()
    push!(values, prev_val)

    while true
        xs_iter = iterate(it.xs, xs_state)

        if xs_iter === nothing
            keep_going = false
            break
        end

        val, xs_state = xs_iter
        key = it.keyfunc(val)

        if key == prev_key
            push!(values, val)
        else
            prev_key = key
            prev_val = val
            break
        end
    end

    return (values, (keep_going, prev_key, prev_val, xs_state))
end