adds the nth function for iterables (#56580)

Hi,

I've turned the open ended issue #54454 into an actual PR.
Tangentially related to #10092 ?

This PR introduces the `nth(itr, n)` function to iterators to give a
`getindex` type of behaviour.
I've tried my best to optimize as much as possible by specializing on
different types of iterators.
In the spirit of iterators any OOB access returns `nothing`. (edit:
instead of throwing an error, i.e. `first(itr, n)` and `last(itr, n)`)

here is the comparison of running the testsuite (~22 different
iterators) using generic `nth` and specialized `nth`:
```julia
@Btime begin                                                                                                                                                                                                                     
    for (itr, n, _) in $testset                                                                                                                                                                                           
         _fallback_nth(itr, n)                                                                                                                                                                                                           
    end                                                                                                                                                                                                                          
end                                                                                                                                                                                                                              
117.750 μs (366 allocations: 17.88 KiB)

@Btime begin                                                                                                                                                                                                                     
  for (itr, n, _) in $testset                                                                                                                                                                                           
    nth(itr, n)                                                                                                                                                                                                              
  end                                                                                                                                                                                                                          
end                                                                                                                                                                                                                              
24.250 μs (341 allocations: 16.70 KiB)
```

---------

Co-authored-by: adienes <51664769+adienes@users.noreply.github.com>
Co-authored-by: Steven G. Johnson <stevenj@mit.edu>
Co-authored-by: Dilum Aluthge <dilum@aluthge.com>

Author: 4 people

NEWS.md

@@ -5,6 +5,7 @@ New language features
 ---------------------
 
   - New `Base.@acquire` macro for a non-closure version of `Base.acquire(f, s::Base.Semaphore)`, like `@lock`. ([#56845])
+  - New `nth` function to access the `n`-th element of a generic iterable. ([#56580])
   - The character U+1F8B2 🢲 (RIGHTWARDS ARROW WITH LOWER HOOK), newly added by Unicode 16,
     is now a valid operator with arrow precedence, accessible as `\hookunderrightarrow` at the REPL.
     ([JuliaLang/JuliaSyntax.jl#525], [#57143])

base/iterators.jl

@@ -16,7 +16,7 @@ using .Base:
     (:), |, +, -, *, !==, !, ==, !=, <=, <, >, >=, =>, missing,
     any, _counttuple, eachindex, ntuple, zero, prod, reduce, in, firstindex, lastindex,
     tail, fieldtypes, min, max, minimum, zero, oneunit, promote, promote_shape, LazyString,
-    afoldl
+    afoldl, mod1
 using .Core
 using Core: @doc
 
@@ -32,7 +32,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
 
 """
@@ -991,6 +991,7 @@ end
 reverse(it::Cycle) = Cycle(reverse(it.xs))
 last(it::Cycle) = last(it.xs)
 
+
 # Repeated - repeat an object infinitely many times
 
 struct Repeated{O}
@@ -1602,4 +1603,83 @@ 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. Throw a `BoundsError`[@ref] if not existing.
+Will advance any `Stateful`[@ref] iterator.
+
+See also: [`first`](@ref), [`last`](@ref)
+
+# Examples
+```jldoctest
+julia> Iterators.nth(2:2:10, 4)
+8
+
+julia> Iterators.nth(reshape(1:30, (5,6)), 6)
+6
+
+julia> stateful = Iterators.Stateful(1:10); Iterators.nth(stateful, 7)
+7
+
+julia> first(stateful)
+8
+```
+"""
+nth(itr, n::Integer) = _nth(IteratorSize(itr), itr, n)
+nth(itr::Cycle{I}, n::Integer) where I = _nth(IteratorSize(I), itr, n)
+nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where O = _nth(IteratorSize(O), itr, n)
+@propagate_inbounds nth(itr::AbstractArray, n::Integer) = itr[begin + n - 1]
+
+function _nth(::Union{HasShape, HasLength}, itr::Cycle{I}, n::Integer) where {I}
+    N = length(itr.xs)
+    N == 0 && throw(BoundsError(itr, n))
+
+    # prevents wrap around behaviour and inherit the error handling
+    return nth(itr.xs, n > 0 ? mod1(n, N) : n)
+end
+
+# Flatten{Take{Repeated{O}}} is the actual type of an Iterators.cycle(iterable::O, m) iterator
+function _nth(::Union{HasShape, HasLength}, itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O}
+    cycles = itr.it.n
+    torepeat = itr.it.xs.x
+    k = length(torepeat)
+    (n > k*cycles || k == 0) && throw(BoundsError(itr, n))
+
+    # prevent wrap around behaviour and inherit the error handling
+    return nth(torepeat, n > 0 ? mod1(n, k) : n)
+end
+
+function _nth(::IteratorSize, itr, n::Integer)
+    # unrolled version of `first(drop)`
+    n > 0 || throw(BoundsError(itr, n))
+    y = iterate(itr)
+    for _ in 1:n-1
+        y === nothing && break
+        y = iterate(itr, y[2])
+    end
+    y === nothing && throw(BoundsError(itr, n))
+    y[1]
+end
+"""
+    nth(n::Integer)
+
+Return a function that gets the `n`-th element from any iterator passed to it.
+Equivalent to `Base.Fix2(nth, n)` or `itr -> nth(itr, n)`.
+
+See also: [`nth`](@ref), [`Base.Fix2`](@ref)
+# Examples
+```jldoctest
+julia> fifth_element = Iterators.nth(5)
+(::Base.Fix2{typeof(Base.Iterators.nth), Int64}) (generic function with 2 methods)
+
+julia> fifth_element(reshape(1:30, (5,6)))
+5
+
+julia> map(fifth_element, ("Willis", "Jovovich", "Oldman"))
+('i', 'v', 'a')
+```
+"""
+nth(n::Integer) = Base.Fix2(nth, n)
+
 end

test/iterators.jl

@@ -1139,6 +1139,65 @@ end
     end
 end
 
+@testset "nth" begin
+
+    Z = Array{Int,0}(undef)
+    Z[] = 17
+    it_result_pairs = Dict(
+        (Z, 1) => 17,
+        (collect(1:100), 23) => 23,
+        (10:6:1000, 123) => 10 + 6 * 122,
+        ("∀ϵ>0", 3) => '>',
+        ((1, 3, 5, 10, 78), 2) => 3,
+        (reshape(1:30, (5, 6)), 21) => 21,
+        (3, 1) => 3,
+        (true, 1) => true,
+        ('x', 1) => 'x',
+        (4 => 5, 2) => 5,
+        (view(Z), 1) => 17,
+        (view(reshape(1:30, (5, 6)), 2:4, 2:6), 10) => 22,
+        ((x^2 for x in 1:10), 9) => 81,
+        (Iterators.Filter(isodd, 1:10), 3) => 5,
+        (Iterators.flatten((1:10, 50:60)), 15) => 54,
+        (pairs(50:60), 7) => 7 => 56,
+        (zip(1:10, 21:30, 51:60), 6) => (6, 26, 56),
+        (Iterators.product(1:3, 10:12), 3) => (3, 10),
+        (Iterators.repeated(3.14159, 5), 4) => 3.14159,
+        ((a=2, b=3, c=5, d=7, e=11), 4) => 7,
+        (Iterators.cycle(collect(1:100)), 9999) => 99,
+        (Iterators.cycle([1, 2, 3, 4, 5], 5), 25) => 5,
+        (Iterators.cycle("String", 10), 16) => 'i',
+        (Iterators.cycle(((),)), 1000) => ()
+    )
+
+
+    @testset "iter: $IT" for (IT, n) in keys(it_result_pairs)
+        @test it_result_pairs[(IT, n)] == nth(IT, n)
+        @test_throws BoundsError nth(IT, -42)
+
+        IT isa Iterators.Cycle && continue # cycles are infinite so never OOB
+        @test_throws BoundsError nth(IT, 999999999)
+    end
+
+    empty_cycle = Iterators.cycle([])
+    @test_throws BoundsError nth(empty_cycle, 42)
+
+    # test the size unknown branch for cycles
+    # only generate odd numbers so we know the actual length
+    # but the iterator is still SizeUnknown()
+    it_size_unknown = Iterators.filter(isodd, 1:2:10)
+    @test Base.IteratorSize(it_size_unknown) isa Base.SizeUnknown
+    @test length(collect(it_size_unknown)) == 5
+
+    cycle_size_unknown = Iterators.cycle(it_size_unknown)
+    finite_cycle_size_unknown = Iterators.cycle(it_size_unknown, 5)
+    @test nth(cycle_size_unknown, 2) == 3
+    @test nth(cycle_size_unknown, 20) == 9 # mod1(20, 5) = 5, wraps 4 times
+    @test nth(finite_cycle_size_unknown, 2) == 3
+    @test nth(finite_cycle_size_unknown, 20) == 9
+    @test_throws BoundsError nth(finite_cycle_size_unknown, 30) # only wraps 5 times, max n is 5 * 5 = 25
+end
+
 @testset "Iterators docstrings" begin
     @test isempty(Docs.undocumented_names(Iterators))
 end