Repository: JuliaLang/julia
File: base/iterators.jl (Julia)

Code before:
julia> first(stateful)
8
```
"""
nth(itr, n::Integer) = _nth(itr, n)

# Count
nth(itr::Count, n::Integer) = n > 0 ? itr.start + itr.step * (n - 1) : throw(ArgumentError("n must be positive."))
# Repeated
nth(itr::Repeated, ::Integer) = itr.x
# Take(Repeated)
nth(itr::Take{Repeated{T}}, n::Integer) where T = begin
    n > itr.n ? throw(BoundsError(itr, n)) : nth(itr.xs, n)
end

# infinite cycle
nth(itr::Cycle{I}, n::Integer) where {I} = begin
    if IteratorSize(I) isa Union{HasShape, HasLength}
        _nth(itr.xs, mod1(n, length(itr.xs)))
    else
        _nth(itr, n)
    end
end

# finite cycle: in reality a Flatten{Take{Repeated{O}}} iterator
nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O} = begin
    if IteratorSize(O) isa Union{HasShape, HasLength}
        cycles = itr.it.n
        repeated = itr.it.xs.x
        k = length(repeated)
        n > k*cycles ? throw(BoundsError(itr, n)) : _nth(repeated, mod1(n, k))
    else
        _nth(itr, n)
    end
end

@inline _nth(itr, n) = first(drop(itr, n-1))
@inline _nth(itr::AbstractArray, n) = itr[begin + n-1]
"""
    nth(n::Integer)

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)
# Examples
```jldoctest
julia> fifth_element = Iterators.nth(5)

Review comment (suggestion):
```suggestion
function nth(itr::Flatten{Take{Repeated{O}}}, n::Integer) where {O}
```

Code after:

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)