Add the take_while, drop_while, take_last and drop_last operators

Ports Enumerable.TakeWhile, SkipWhile, TakeLast and SkipLast. The Skip
operators are named drop_* to match the existing `drop` rather than LINQ's
`Skip`.

take_while and drop_while are lazy: take_while stops at the first element
that fails the predicate and never examines the rest, and drop_while passes
everything through untested once the leading run is gone.

take_last holds at most n elements in a CircularBuffer while walking the
source. drop_last stays n elements behind the source, emitting an element
only once n further ones have been read.

A count of zero or less yields nothing for take_last and leaves the source
unchanged for drop_last, matching .NET.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: davidanthoff

src/QueryOperators.jl

@@ -26,6 +26,10 @@ include("enumerable/enumerable_defaultifempty.jl")
 include("enumerable/enumerable_count.jl")
 include("enumerable/enumerable_take.jl")
 include("enumerable/enumerable_drop.jl")
+include("enumerable/enumerable_takewhile.jl")
+include("enumerable/enumerable_dropwhile.jl")
+include("enumerable/enumerable_takelast.jl")
+include("enumerable/enumerable_droplast.jl")
 include("enumerable/enumerable_unique.jl")
 include("enumerable/enumerable_concat.jl")
 include("enumerable/enumerable_union.jl")

src/enumerable/enumerable_droplast.jl

@@ -0,0 +1,48 @@
+struct EnumerableDropLast{T,S} <: Enumerable
+    source::S
+    n::Int
+end
+
+# Enumerable.SkipLast: everything but the trailing `n` elements. A count of
+# zero or less leaves the source unchanged, as in .NET.
+function drop_last(source::Enumerable, n::Integer)
+    T = eltype(source)
+    return EnumerableDropLast{T,typeof(source)}(source, max(Int(n), 0))
+end
+
+Base.IteratorSize(::Type{EnumerableDropLast{T,S}}) where {T,S} = haslength(S)
+
+Base.eltype(::Type{EnumerableDropLast{T,S}}) where {T,S} = T
+
+Base.length(iter::EnumerableDropLast) = max(length(iter.source) - iter.n, 0)
+
+# Stays `n` elements behind the source: an element is only emitted once `n`
+# further elements have been read, which proves it is not one of the last `n`.
+function Base.iterate(iter::EnumerableDropLast{T,S}) where {T,S}
+    buffer = T[]
+    source_state = _NotStarted()
+
+    while length(buffer) < iter.n
+        ret = _iterate_from(iter.source, source_state)
+        # Fewer than n elements in total, so every one of them is dropped.
+        ret === nothing && return nothing
+        push!(buffer, ret[1])
+        source_state = ret[2]
+    end
+
+    return _drop_last_next(iter, buffer, source_state)
+end
+
+function Base.iterate(iter::EnumerableDropLast, state)
+    return _drop_last_next(iter, state.buffer, state.state)
+end
+
+function _drop_last_next(iter::EnumerableDropLast, buffer, source_state)
+    ret = _iterate_from(iter.source, source_state)
+    ret === nothing && return nothing
+
+    push!(buffer, ret[1])
+    element = popfirst!(buffer)
+
+    return element, (buffer=buffer, state=ret[2])
+end

src/enumerable/enumerable_dropwhile.jl

@@ -0,0 +1,27 @@
+struct EnumerableDropWhile{T,S,Q<:Function} <: Enumerable
+    source::S
+    f::Q
+end
+
+# Enumerable.SkipWhile: discards the leading run of elements satisfying the
+# predicate, then yields everything that follows without testing it again.
+function drop_while(source::Enumerable, f::Function, f_expr::Expr)
+    T = eltype(source)
+    return EnumerableDropWhile{T,typeof(source),typeof(f)}(source, f)
+end
+
+Base.eltype(::Type{EnumerableDropWhile{T,S,Q}}) where {T,S,Q} = T
+
+function Base.iterate(iter::EnumerableDropWhile)
+    ret = iterate(iter.source)
+
+    while ret !== nothing && iter.f(ret[1])
+        ret = iterate(iter.source, ret[2])
+    end
+
+    ret === nothing && return nothing
+
+    return ret[1], ret[2]
+end
+
+Base.iterate(iter::EnumerableDropWhile, state) = iterate(iter.source, state)

src/enumerable/enumerable_takelast.jl

@@ -0,0 +1,42 @@
+struct EnumerableTakeLast{T,S} <: Enumerable
+    source::S
+    n::Int
+end
+
+# Enumerable.TakeLast: the trailing `n` elements. A count of zero or less
+# yields nothing, as in .NET.
+function take_last(source::Enumerable, n::Integer)
+    T = eltype(source)
+    return EnumerableTakeLast{T,typeof(source)}(source, max(Int(n), 0))
+end
+
+Base.IteratorSize(::Type{EnumerableTakeLast{T,S}}) where {T,S} = haslength(S)
+
+Base.eltype(::Type{EnumerableTakeLast{T,S}}) where {T,S} = T
+
+Base.length(iter::EnumerableTakeLast) = min(length(iter.source), iter.n)
+
+# Which elements are last is only known once the source is exhausted, so the
+# whole source is walked, holding at most `n` elements in a ring buffer.
+function Base.iterate(iter::EnumerableTakeLast{T,S}) where {T,S}
+    iter.n == 0 && return nothing
+
+    buffer = CircularBuffer{T}(iter.n)
+    for i in iter.source
+        push!(buffer, i)
+    end
+
+    length(buffer)==0 && return nothing
+
+    elements = Base.collect(buffer)
+
+    return elements[1], (elements, 2)
+end
+
+function Base.iterate(iter::EnumerableTakeLast{T,S}, state) where {T,S}
+    if state[2]>length(state[1])
+        return nothing
+    else
+        return state[1][state[2]], (state[1], state[2]+1)
+    end
+end

src/enumerable/enumerable_takewhile.jl

@@ -0,0 +1,23 @@
+struct EnumerableTakeWhile{T,S,Q<:Function} <: Enumerable
+    source::S
+    f::Q
+end
+
+# Enumerable.TakeWhile: yields elements until the predicate first fails, then
+# stops — later elements are not examined even if they would satisfy it.
+function take_while(source::Enumerable, f::Function, f_expr::Expr)
+    T = eltype(source)
+    return EnumerableTakeWhile{T,typeof(source),typeof(f)}(source, f)
+end
+
+Base.eltype(::Type{EnumerableTakeWhile{T,S,Q}}) where {T,S,Q} = T
+
+Base.iterate(iter::EnumerableTakeWhile) = _take_while_step(iter, iterate(iter.source))
+
+Base.iterate(iter::EnumerableTakeWhile, state) = _take_while_step(iter, iterate(iter.source, state))
+
+function _take_while_step(iter::EnumerableTakeWhile, ret)
+    ret === nothing && return nothing
+    iter.f(ret[1]) || return nothing
+    return ret[1], ret[2]
+end

src/operators.jl

@@ -180,3 +180,14 @@ function count_by end
 function aggregate_by end
 
 function chunk end
+
+# Partitioning. `drop_while` and `drop_last` are Enumerable.SkipWhile and
+# SkipLast, named to match the existing `drop` rather than LINQ's `Skip`.
+
+function take_while end
+
+function drop_while end
+
+function take_last end
+
+function drop_last end

test/runtests.jl

@@ -5,6 +5,7 @@ include("test_enumerable_outerjoins.jl")
 include("test_enumerable_setops.jl")
 include("test_enumerable_ordering.jl")
 include("test_enumerable_keyed_aggregation.jl")
+include("test_enumerable_partitioning.jl")
 include("test_enumerable_unique.jl")
 include("test_enumerable_summarize.jl")
 include("test_namedtupleutilities.jl")

test/test_enumerable_partitioning.jl

@@ -0,0 +1,165 @@
+@testitem "take_while" begin
+    source = QueryOperators.query([1, 2, 3, 4, 1, 2])
+
+    res = QueryOperators.take_while(source, i -> i < 3, :(i -> i < 3))
+
+    # Stops at the first failure; the trailing 1 and 2 are not taken.
+    @test collect(res) == [1, 2]
+    @test eltype(res) == Int
+end
+
+@testitem "take_while taking everything or nothing" begin
+    source = QueryOperators.query([1, 2, 3])
+
+    @test collect(QueryOperators.take_while(source, i -> true, :(i -> true))) == [1, 2, 3]
+    @test collect(QueryOperators.take_while(source, i -> false, :(i -> false))) == Int[]
+end
+
+@testitem "take_while on an empty source" begin
+    source = QueryOperators.query(Int[])
+
+    @test collect(QueryOperators.take_while(source, i -> true, :(i -> true))) == Int[]
+end
+
+@testitem "take_while stops walking the source at the first failure" begin
+    import IteratorInterfaceExtensions
+
+    mutable struct CountedTakeWhileSource
+        data::Vector{Int}
+        pulled::Int
+    end
+    Base.eltype(::Type{CountedTakeWhileSource}) = Int
+    Base.IteratorSize(::Type{CountedTakeWhileSource}) = Base.HasLength()
+    Base.length(c::CountedTakeWhileSource) = length(c.data)
+    function Base.iterate(c::CountedTakeWhileSource, i=1)
+        i > length(c.data) && return nothing
+        c.pulled += 1
+        return c.data[i], i + 1
+    end
+    IteratorInterfaceExtensions.isiterable(::CountedTakeWhileSource) = true
+    IteratorInterfaceExtensions.getiterator(c::CountedTakeWhileSource) = c
+
+    src = CountedTakeWhileSource(collect(1:100), 0)
+    res = QueryOperators.take_while(QueryOperators.query(src), i -> i < 4, :(i -> i < 4))
+
+    @test collect(res) == [1, 2, 3]
+    # Three taken plus the one that failed the predicate.
+    @test src.pulled == 4
+end
+
+@testitem "drop_while" begin
+    source = QueryOperators.query([1, 2, 3, 4, 1, 2])
+
+    res = QueryOperators.drop_while(source, i -> i < 3, :(i -> i < 3))
+
+    # Only the leading run is dropped; the trailing 1 and 2 survive.
+    @test collect(res) == [3, 4, 1, 2]
+    @test eltype(res) == Int
+end
+
+@testitem "drop_while dropping everything or nothing" begin
+    source = QueryOperators.query([1, 2, 3])
+
+    @test collect(QueryOperators.drop_while(source, i -> true, :(i -> true))) == Int[]
+    @test collect(QueryOperators.drop_while(source, i -> false, :(i -> false))) == [1, 2, 3]
+end
+
+@testitem "drop_while on an empty source" begin
+    source = QueryOperators.query(Int[])
+
+    @test collect(QueryOperators.drop_while(source, i -> true, :(i -> true))) == Int[]
+end
+
+@testitem "take_while and drop_while partition the source" begin
+    source = QueryOperators.query([1, 2, 3, 4, 5])
+    p = i -> i < 3
+
+    taken = collect(QueryOperators.take_while(source, p, :(i -> i < 3)))
+    dropped = collect(QueryOperators.drop_while(source, p, :(i -> i < 3)))
+
+    @test vcat(taken, dropped) == [1, 2, 3, 4, 5]
+end
+
+@testitem "take_last" begin
+    source = QueryOperators.query([1, 2, 3, 4, 5])
+
+    res = QueryOperators.take_last(source, 2)
+
+    @test collect(res) == [4, 5]
+    @test eltype(res) == Int
+    @test Base.IteratorSize(typeof(res)) == Base.HasLength()
+    @test length(res) == 2
+end
+
+@testitem "take_last with a count at or beyond the source length" begin
+    source = QueryOperators.query([1, 2, 3])
+
+    @test collect(QueryOperators.take_last(source, 3)) == [1, 2, 3]
+    @test collect(QueryOperators.take_last(source, 10)) == [1, 2, 3]
+    @test length(QueryOperators.take_last(source, 10)) == 3
+end
+
+@testitem "take_last with a count of zero or less" begin
+    source = QueryOperators.query([1, 2, 3])
+
+    @test collect(QueryOperators.take_last(source, 0)) == Int[]
+    @test collect(QueryOperators.take_last(source, -1)) == Int[]
+end
+
+@testitem "take_last on an empty source" begin
+    source = QueryOperators.query(Int[])
+
+    @test collect(QueryOperators.take_last(source, 2)) == Int[]
+end
+
+@testitem "drop_last" begin
+    source = QueryOperators.query([1, 2, 3, 4, 5])
+
+    res = QueryOperators.drop_last(source, 2)
+
+    @test collect(res) == [1, 2, 3]
+    @test eltype(res) == Int
+    @test Base.IteratorSize(typeof(res)) == Base.HasLength()
+    @test length(res) == 3
+end
+
+@testitem "drop_last with a count at or beyond the source length" begin
+    source = QueryOperators.query([1, 2, 3])
+
+    @test collect(QueryOperators.drop_last(source, 3)) == Int[]
+    @test collect(QueryOperators.drop_last(source, 10)) == Int[]
+    @test length(QueryOperators.drop_last(source, 10)) == 0
+end
+
+@testitem "drop_last with a count of zero or less is the identity" begin
+    source = QueryOperators.query([1, 2, 3])
+
+    @test collect(QueryOperators.drop_last(source, 0)) == [1, 2, 3]
+    @test collect(QueryOperators.drop_last(source, -1)) == [1, 2, 3]
+end
+
+@testitem "drop_last on an empty source" begin
+    source = QueryOperators.query(Int[])
+
+    @test collect(QueryOperators.drop_last(source, 2)) == Int[]
+end
+
+@testitem "take_last and drop_last partition the source" begin
+    source = QueryOperators.query([1, 2, 3, 4, 5])
+
+    dropped = collect(QueryOperators.drop_last(source, 2))
+    taken = collect(QueryOperators.take_last(source, 2))
+
+    @test vcat(dropped, taken) == [1, 2, 3, 4, 5]
+end
+
+@testitem "partitioning operators work downstream of groupby" begin
+    source = QueryOperators.query([(k=1, v=1), (k=2, v=2), (k=3, v=3)])
+    grouped = QueryOperators.@groupby_simple(source, i -> i.k)
+
+    last_two = collect(QueryOperators.take_last(grouped, 2))
+    @test [QueryOperators.key(g) for g in last_two] == [2, 3]
+
+    all_but_last = collect(QueryOperators.drop_last(grouped, 1))
+    @test [QueryOperators.key(g) for g in all_but_last] == [1, 2]
+end