From 6a0b4bd78ca389ae5d0c4604e76062cca66e1c9b Mon Sep 17 00:00:00 2001 From: cschen Date: Sat, 16 Nov 2024 13:00:25 +0100 Subject: [PATCH 1/6] adds `nth` function to iterators plus tests --- base/iterators.jl | 71 +++++++++++++++++++++++++++++++++++++++++++++++ test/iterators.jl | 24 ++++++++++++++++ 2 files changed, 95 insertions(+) diff --git a/base/iterators.jl b/base/iterators.jl index 6b8d9fe75e302..09f8c6edcacc7 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1595,4 +1595,75 @@ end # be the same as the keys, so this is a valid optimization (see #51631) pairs(s::AbstractString) = IterableStatePairs(s) +""" + nth(itr, n::Integer) + +Get the `n`th element of an iterable collection. Return `nothing` if not existing. + +See also: [`first`](@ref), [`last`](@ref) + +# Examples +```jldoctest +julia> nth(2:2:10, 4) +8 + +julia> nth(reshape(1:30, (5,6)), 6) +6 +``` +""" +nth(itr, n::Integer) = _nth(IteratorSize(itr), itr, n) +nth(itr::AbstractArray, n::Integer) = n > length(itr) ? nothing : itr[n] + +_nth(::SizeUnknown, itr, n) = _fallback_nth(itr, n) +_nth(::Union{HasShape,HasLength}, itr, n) = _withlength_nth(itr, n, length(itr)) +_nth(::IsInfinite, itr, n) = _inbounds_nth(itr, n) + +_inbounds_nth(itr, n) = iterate(Base.Iterators.drop(itr, n - 1))[1] +_inbounds_nth(itr::AbstractArray, n) = itr[n] + +_withlength_nth(itr, n, N) = n > N ? nothing : _inbounds_nth(itr, n) + +function _fallback_nth(itr, n) + y = iterate(Iterators.drop(itr, n - 1)) + y === nothing && return nothing + y[1] +end + +# specialized versions to better interact with existing iterators +# Count +nth(itr::Count, n::Integer) = n > 0 ? itr.start + itr.step * (n - 1) : nothing + +# Repeated +nth(itr::Repeated, n::Integer) = itr.x + +# Take(Repeated) +nth(itr::Take{Repeated{O}}, n::Integer) where {O} = n > itr.n ? nothing : itr.xs.x + +# infinite cycle +nth(itr::Iterators.Cycle{I}, n::Integer) where {I} = _nth_inf_cycle(IteratorSize(I), itr, n) +_nth_inf_cycle(::IsInfinite, itr, n) = _inbounds_nth(itr.xs, n) +_nth_inf_cycle(::SizeUnknown, itr, n) = _fallback_nth(itr.xs, n) +_nth_inf_cycle(::Union{HasShape,HasLength}, itr, n) = _repeating_cycle_nth(itr.xs, n, length(itr.xs)) + +# finite cycle +# a finite cycle iterator is in reality a Flatten{Take{Repeated{O}}} iterator +nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O} = _nth_finite_cycle(IteratorSize(O), itr, n) +_nth_finite_cycle(::IsInfinite, itr, n) = _inbounds_nth(itr, n) +_nth_finite_cycle(::SizeUnknown, itr, n) = _fallback_nth(itr, n) +_nth_finite_cycle(::Union{HasShape,HasLength}, itr, n) = begin + N = itr.it.n # `Take` iterator n + torepeat = itr.it.xs.x # repeated object + K = length(torepeat) + n > K * N && return nothing + _repeating_cycle_nth(torepeat, n, K) +end +# O = [....] # K length +# Repeat(O) = [ [....] [....] [....] ... +# Take(Repeat(O), N) = [ [....] [....] [....] ] # N times +# Flatten(...) = [....|....|....] # K*N elements + +_repeating_cycle_nth(inner_itr, n, inner_N) = _inbounds_nth(inner_itr, 1 + ((n - 1) % inner_N)) + + + end diff --git a/test/iterators.jl b/test/iterators.jl index d1e7525c43465..aabb39d5cc7ba 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -1074,6 +1074,30 @@ end end end +@testset "nth" begin + Z = Array{Int,0}(undef) + Z[] = 17 + itrs = (collect(1:1000), 10:6:1000, "∀ϵ>0", (1, 3, 5, 10, 78), reshape(1:30, (5, 6)), + Z, 3, true, 'x', 4 => 5, view(Z), view(reshape(1:30, (5, 6)), 2:4, 2:6), + (x^2 for x in 1:10), Iterators.Filter(isodd, 1:10), Iterators.flatten((1:10, 50:60)), + pairs(50:60), zip(1:10, 21:30, 51:60), Iterators.product(1:3, 10:12), + Iterators.repeated(3.14159, 5), (a=2, b=3, c=5, d=7, e=11), Iterators.cycle(collect(1:100)), + Iterators.cycle([1, 2, 3, 4, 5], 5)) + ns = ( + 234, 123, 3, 2, 21, 1, 1, 1, 1, 1, 1, 10, 9, 3, 15, 7, 6, 3, 4, 4, 99999, 25 + ) + expected = ( + 234, 742, '>', 3, 21, 17, 3, true, 'x', 4, 17, 22, 81, 5, 54, (7 => 56), (6, 26, 56), (3, 10), 3.14159, 7, 99, 5 + ) + @test length(itrs) == length(ns) == length(expected) + testset = zip(itrs, ns, expected) + @testset "iter: $IT" for (IT, n, exp) in testset + @test exp == nth(IT, n) + IT isa Cycle && continue # cycles are infinite so never OOB + @test nth(IT, 999999999) === nothing + end +end + @testset "Iterators docstrings" begin @test isempty(Docs.undocumented_names(Iterators)) end From 1de4a8fe15def99d2112c28173d4c42c71db8a0d Mon Sep 17 00:00:00 2001 From: cschen Date: Sat, 16 Nov 2024 13:21:32 +0100 Subject: [PATCH 2/6] removed useless module annotations added nth to export list --- base/iterators.jl | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index 09f8c6edcacc7..af78dd69bf4f9 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -30,7 +30,7 @@ import .Base: getindex, setindex!, get, iterate, popfirst!, isdone, peek, intersect -export enumerate, zip, rest, countfrom, take, drop, takewhile, dropwhile, cycle, repeated, product, flatten, flatmap, partition +export enumerate, zip, rest, countfrom, take, drop, takewhile, dropwhile, cycle, repeated, product, flatten, flatmap, partition, nth public accumulate, filter, map, peel, reverse, Stateful """ @@ -1618,13 +1618,13 @@ _nth(::SizeUnknown, itr, n) = _fallback_nth(itr, n) _nth(::Union{HasShape,HasLength}, itr, n) = _withlength_nth(itr, n, length(itr)) _nth(::IsInfinite, itr, n) = _inbounds_nth(itr, n) -_inbounds_nth(itr, n) = iterate(Base.Iterators.drop(itr, n - 1))[1] +_inbounds_nth(itr, n) = iterate(drop(itr, n - 1))[1] _inbounds_nth(itr::AbstractArray, n) = itr[n] _withlength_nth(itr, n, N) = n > N ? nothing : _inbounds_nth(itr, n) function _fallback_nth(itr, n) - y = iterate(Iterators.drop(itr, n - 1)) + y = iterate(drop(itr, n - 1)) y === nothing && return nothing y[1] end @@ -1640,7 +1640,7 @@ nth(itr::Repeated, n::Integer) = itr.x nth(itr::Take{Repeated{O}}, n::Integer) where {O} = n > itr.n ? nothing : itr.xs.x # infinite cycle -nth(itr::Iterators.Cycle{I}, n::Integer) where {I} = _nth_inf_cycle(IteratorSize(I), itr, n) +nth(itr::Cycle{I}, n::Integer) where {I} = _nth_inf_cycle(IteratorSize(I), itr, n) _nth_inf_cycle(::IsInfinite, itr, n) = _inbounds_nth(itr.xs, n) _nth_inf_cycle(::SizeUnknown, itr, n) = _fallback_nth(itr.xs, n) _nth_inf_cycle(::Union{HasShape,HasLength}, itr, n) = _repeating_cycle_nth(itr.xs, n, length(itr.xs)) @@ -1657,10 +1657,7 @@ _nth_finite_cycle(::Union{HasShape,HasLength}, itr, n) = begin n > K * N && return nothing _repeating_cycle_nth(torepeat, n, K) end -# O = [....] # K length -# Repeat(O) = [ [....] [....] [....] ... -# Take(Repeat(O), N) = [ [....] [....] [....] ] # N times -# Flatten(...) = [....|....|....] # K*N elements + _repeating_cycle_nth(inner_itr, n, inner_N) = _inbounds_nth(inner_itr, 1 + ((n - 1) % inner_N)) From 3870e3ea7aeb2f6b7c7a23608d4dcda0af057a8b Mon Sep 17 00:00:00 2001 From: cschen Date: Sat, 16 Nov 2024 15:18:37 +0100 Subject: [PATCH 3/6] fix one based indexing, more generic add docs explaining interaction with Stateful iterators change test to be Any vectors instead of tuples (actually way faster as well) --- base/iterators.jl | 12 ++++++++++-- test/iterators.jl | 16 ++++++++-------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index af78dd69bf4f9..951e496b40076 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1599,6 +1599,7 @@ pairs(s::AbstractString) = IterableStatePairs(s) nth(itr, n::Integer) Get the `n`th element of an iterable collection. Return `nothing` if not existing. +Will advance any `Stateful`[@ref] iterator. See also: [`first`](@ref), [`last`](@ref) @@ -1609,17 +1610,24 @@ julia> nth(2:2:10, 4) julia> nth(reshape(1:30, (5,6)), 6) 6 + +julia> stateful = Stateful(1:10); nth(stateful, 7) +7 + +julia> first(stateful) +8 ``` """ nth(itr, n::Integer) = _nth(IteratorSize(itr), itr, n) -nth(itr::AbstractArray, n::Integer) = n > length(itr) ? nothing : itr[n] +nth(itr::AbstractArray, n::Integer) = n > length(itr) ? nothing : _inbounds_nth(itr, n) +nth(itr::AbstractRange, n) = getindex(itr, n) _nth(::SizeUnknown, itr, n) = _fallback_nth(itr, n) _nth(::Union{HasShape,HasLength}, itr, n) = _withlength_nth(itr, n, length(itr)) _nth(::IsInfinite, itr, n) = _inbounds_nth(itr, n) _inbounds_nth(itr, n) = iterate(drop(itr, n - 1))[1] -_inbounds_nth(itr::AbstractArray, n) = itr[n] +_inbounds_nth(itr::AbstractArray, n) = getindex(itr, nth(eachindex(IndexLinear(), itr), n)) _withlength_nth(itr, n, N) = n > N ? nothing : _inbounds_nth(itr, n) diff --git a/test/iterators.jl b/test/iterators.jl index aabb39d5cc7ba..c2aa72198bec6 100644 --- a/test/iterators.jl +++ b/test/iterators.jl @@ -1077,18 +1077,18 @@ end @testset "nth" begin Z = Array{Int,0}(undef) Z[] = 17 - itrs = (collect(1:1000), 10:6:1000, "∀ϵ>0", (1, 3, 5, 10, 78), reshape(1:30, (5, 6)), - Z, 3, true, 'x', 4 => 5, view(Z), view(reshape(1:30, (5, 6)), 2:4, 2:6), + itrs = Any[collect(1:1000), 10:6:1000, "∀ϵ>0", (1, 3, 5, 10, 78), reshape(1:30, (5, 6)), + Z, 3, true, 'x', 4=>5, view(Z), view(reshape(1:30, (5, 6)), 2:4, 2:6), (x^2 for x in 1:10), Iterators.Filter(isodd, 1:10), Iterators.flatten((1:10, 50:60)), pairs(50:60), zip(1:10, 21:30, 51:60), Iterators.product(1:3, 10:12), Iterators.repeated(3.14159, 5), (a=2, b=3, c=5, d=7, e=11), Iterators.cycle(collect(1:100)), - Iterators.cycle([1, 2, 3, 4, 5], 5)) - ns = ( + Iterators.cycle([1, 2, 3, 4, 5], 5)] + ns = Any[ 234, 123, 3, 2, 21, 1, 1, 1, 1, 1, 1, 10, 9, 3, 15, 7, 6, 3, 4, 4, 99999, 25 - ) - expected = ( - 234, 742, '>', 3, 21, 17, 3, true, 'x', 4, 17, 22, 81, 5, 54, (7 => 56), (6, 26, 56), (3, 10), 3.14159, 7, 99, 5 - ) + ] + expected = Any[ + 234, 742, '>', 3, 21, 17, 3, true, 'x', 4, 17, 22, 81, 5, 54, (7=>56), (6, 26, 56), (3, 10), 3.14159, 7, 99, 5 + ] @test length(itrs) == length(ns) == length(expected) testset = zip(itrs, ns, expected) @testset "iter: $IT" for (IT, n, exp) in testset From 8bdd4a089cef5d165374527bf0c6a248f63beff1 Mon Sep 17 00:00:00 2001 From: cschen Date: Sun, 17 Nov 2024 09:51:54 +0100 Subject: [PATCH 4/6] fix: simpler generic-based indexing --- base/iterators.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/base/iterators.jl b/base/iterators.jl index 951e496b40076..585ffe75afadd 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1627,7 +1627,7 @@ _nth(::Union{HasShape,HasLength}, itr, n) = _withlength_nth(itr, n, length(itr)) _nth(::IsInfinite, itr, n) = _inbounds_nth(itr, n) _inbounds_nth(itr, n) = iterate(drop(itr, n - 1))[1] -_inbounds_nth(itr::AbstractArray, n) = getindex(itr, nth(eachindex(IndexLinear(), itr), n)) +_inbounds_nth(itr::AbstractArray, n) = itr[begin + n-1] _withlength_nth(itr, n, N) = n > N ? nothing : _inbounds_nth(itr, n) From 7326e870444c1e5fb22a27066e21cfbddd87a40c Mon Sep 17 00:00:00 2001 From: cschen Date: Sun, 5 Jan 2025 13:46:33 +0100 Subject: [PATCH 5/6] feat: adds the fix2 wrapper. chore: small reordering. feat: nth is only throwing. with explicit errors. --- base/iterators.jl | 77 ++++++++++++++++++++++++++++++----------------- 1 file changed, 49 insertions(+), 28 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index 585ffe75afadd..65aecccb0dc38 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1598,10 +1598,10 @@ pairs(s::AbstractString) = IterableStatePairs(s) """ nth(itr, n::Integer) -Get the `n`th element of an iterable collection. Return `nothing` if not existing. +Get the `n`th element of an iterable collection. Throw a `BoundsError`[@ref] if not existing. Will advance any `Stateful`[@ref] iterator. -See also: [`first`](@ref), [`last`](@ref) +See also: [`first`](@ref), [`last`](@ref), [`nth`](@ref) # Examples ```jldoctest @@ -1619,56 +1619,77 @@ julia> first(stateful) ``` """ nth(itr, n::Integer) = _nth(IteratorSize(itr), itr, n) -nth(itr::AbstractArray, n::Integer) = n > length(itr) ? nothing : _inbounds_nth(itr, n) -nth(itr::AbstractRange, n) = getindex(itr, n) +nth(itr::AbstractArray, n::Integer) = _withlength_nth_throw(itr, n, length(itr)) +# specialized versions to better interact with existing iterators +# Count +nth(itr::Count, n::Integer) = n > 0 ? itr.start + itr.step * (n - 1) : throw(ArgumentError("n must be positive.")) +# # Repeated +nth(itr::Repeated, n::Integer) = itr.x +# # Take(Repeated) +nth(itr::Take{Repeated{O}}, n::Integer) where {O} = begin + n > itr.n ? throw(BoundsError("attempted to access $(itr.n)-element $(typeof(itr)) at position $n")) : itr.xs.x +end +# infinite cycle +nth(itr::Cycle{I}, n::Integer) where {I} = _nth_inf_cycle(IteratorSize(I), itr, n) +# finite cycle: in reality a Flatten{Take{Repeated{O}}} iterator +nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O} = _nth_finite_cycle(IteratorSize(O), itr, n) -_nth(::SizeUnknown, itr, n) = _fallback_nth(itr, n) -_nth(::Union{HasShape,HasLength}, itr, n) = _withlength_nth(itr, n, length(itr)) +_nth(::SizeUnknown, itr, n) = _fallback_nth_throw(itr, n) +_nth(::Union{HasShape,HasLength}, itr, n) = _withlength_nth_throw(itr, n, length(itr)) _nth(::IsInfinite, itr, n) = _inbounds_nth(itr, n) _inbounds_nth(itr, n) = iterate(drop(itr, n - 1))[1] _inbounds_nth(itr::AbstractArray, n) = itr[begin + n-1] -_withlength_nth(itr, n, N) = n > N ? nothing : _inbounds_nth(itr, n) +function _withlength_nth_throw(itr, n, N) + n > N && throw(BoundsError("attempted to access $N-element $(typeof(itr)) at position $n")) + _inbounds_nth(itr, n) +end -function _fallback_nth(itr, n) +function _fallback_nth_throw(itr, n) y = iterate(drop(itr, n - 1)) - y === nothing && return nothing + y === nothing && throw(BoundsError("Iterator $(typeof(itr)) has less than $n elements")) y[1] end -# specialized versions to better interact with existing iterators -# Count -nth(itr::Count, n::Integer) = n > 0 ? itr.start + itr.step * (n - 1) : nothing - -# Repeated -nth(itr::Repeated, n::Integer) = itr.x - -# Take(Repeated) -nth(itr::Take{Repeated{O}}, n::Integer) where {O} = n > itr.n ? nothing : itr.xs.x - -# infinite cycle -nth(itr::Cycle{I}, n::Integer) where {I} = _nth_inf_cycle(IteratorSize(I), itr, n) _nth_inf_cycle(::IsInfinite, itr, n) = _inbounds_nth(itr.xs, n) _nth_inf_cycle(::SizeUnknown, itr, n) = _fallback_nth(itr.xs, n) _nth_inf_cycle(::Union{HasShape,HasLength}, itr, n) = _repeating_cycle_nth(itr.xs, n, length(itr.xs)) -# finite cycle -# a finite cycle iterator is in reality a Flatten{Take{Repeated{O}}} iterator -nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O} = _nth_finite_cycle(IteratorSize(O), itr, n) _nth_finite_cycle(::IsInfinite, itr, n) = _inbounds_nth(itr, n) -_nth_finite_cycle(::SizeUnknown, itr, n) = _fallback_nth(itr, n) -_nth_finite_cycle(::Union{HasShape,HasLength}, itr, n) = begin +_nth_finite_cycle(::SizeUnknown, itr, n) = _fallback_nth_throw(itr, n) +_nth_finite_cycle(::Union{HasShape,HasLength}, itr, n) = _walk_cycle_throw(itr, n) + +_repeating_cycle_nth(inner_itr, n, inner_N) = _inbounds_nth(inner_itr, 1 + ((n - 1) % inner_N)) + +function _walk_cycle_throw(itr, n) N = itr.it.n # `Take` iterator n torepeat = itr.it.xs.x # repeated object K = length(torepeat) - n > K * N && return nothing + n > K * N && throw(BoundsError("attempted to access $(N*K)-element $(typeof(itr)) at position $n")) _repeating_cycle_nth(torepeat, n, K) end +""" + nth(n::Integer) -_repeating_cycle_nth(inner_itr, n, inner_N) = _inbounds_nth(inner_itr, 1 + ((n - 1) % inner_N)) +return a function that gets the `n`-th element from any iterator passed to it. +Fixes the second element with +Throw a `BoundsError`[@ref] if not existing. +Will advance any `Stateful`[@ref] iterator. +See also: [`nth`](@ref), [`Base.Fix2`](@ref) +# Examples +```jldoctest +julia> fifth_element = nth(5) +(::Base.Fix2{typeof(nth), Int64}) (generic function with 1 method) +julia> fifth_element(reshape(1:30, (5,6))) +5 + +julia> map(nth(3), my_vec) +``` +""" +nth(n::Integer) = Base.Fix2(nth, n) end From 58b60e8525fdc0653400bcee4e8f80e64304242b Mon Sep 17 00:00:00 2001 From: cschen Date: Fri, 10 Jan 2025 12:58:19 +0100 Subject: [PATCH 6/6] fix: better wording for the fix2 version --- base/iterators.jl | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/base/iterators.jl b/base/iterators.jl index 65aecccb0dc38..e1ad78d9d7f64 100644 --- a/base/iterators.jl +++ b/base/iterators.jl @@ -1673,9 +1673,10 @@ end """ nth(n::Integer) -return a function that gets the `n`-th element from any iterator passed to it. -Fixes the second element with +Return a function that gets the `n`-th element from any iterator passed to it. Throw a `BoundsError`[@ref] if not existing. + +Fixes the second element. Equivalent to `Base.Fix2(nth, n)`. Will advance any `Stateful`[@ref] iterator. See also: [`nth`](@ref), [`Base.Fix2`](@ref)