ronantakizawa/codereview-bench huggingface dataset:
Repository: JuliaLang/julia
File: base/iterators.jl (Julia)

Code before:

"""
    findeach(f, it)
    findeach(it)

An iterator that generates every key from the key/value pairs of `pairs(it)`,
where `f(value)` returns `true`.

If `f` is not specified, default to `identity`.

`Iterators.findeach` is the lazy equivalent of `findall`.

!!! compat "Julia 1.13"
    Lazy `findeach` requires at least Julia 1.13.

# Examples
```jldoctest
julia> collect(Iterators.findeach(isodd, Dict(2 => 3, 3 => 2)))
1-element Vector{Int64}:
 2

julia> only(Iterators.findeach(==(1), [3,6,2,1]))
4
```
"""
findeach(f, it) = map(first, filter(i -> f(last(i)), pairs(it)))

findeach(it) = findeach(identity, it)


# Cycle an iterator forever

struct Cycle{I}
    xs::I
end

"""
    cycle(iter[, n::Int])

An iterator that cycles through `iter` forever.
If `n` is specified, then it cycles through `iter` that many times.
When `iter` is empty, so are `cycle(iter)` and `cycle(iter, n)`.

`Iterators.cycle(iter, n)` is the lazy equivalent of [`Base.repeat`](@ref)`(vector, n)`,
while [`Iterators.repeated`](@ref)`(iter, n)` is the lazy [`Base.fill`](@ref)`(item, n)`.

!!! compat "Julia 1.11"
    The method `cycle(iter, n)` was added in Julia 1.11.

# Examples
```jldoctest

Review comment (suggestion):
Could also write this a bit more succinctly as
```suggestion
findeach(f, it) = (k for (k, v) in pairs(it) if f(v))
```

Code after:

"""
    findeach(f, it)
    findeach(it)

An iterator that generates every key from the key/value pairs of `pairs(it)`,
where `f(value)` returns `true`.

If `f` is not specified, default to `identity`.

`Iterators.findeach` is the lazy equivalent of `findall`.

!!! compat "Julia 1.13"
    `findeach` requires at least Julia 1.13.

# Examples
```jldoctest
julia> collect(Iterators.findeach(isodd, Dict(2 => 3, 3 => 2)))
1-element Vector{Int64}:
 2

julia> only(Iterators.findeach(==(1), [3,6,2,1]))
4
```
"""
findeach(f, it) = (k for (k, v) in pairs(it) if f(v))

findeach(it) = findeach(identity, it)


# Cycle an iterator forever

struct Cycle{I}
    xs::I
end

"""
    cycle(iter[, n::Int])

An iterator that cycles through `iter` forever.
If `n` is specified, then it cycles through `iter` that many times.
When `iter` is empty, so are `cycle(iter)` and `cycle(iter, n)`.

`Iterators.cycle(iter, n)` is the lazy equivalent of [`Base.repeat`](@ref)`(vector, n)`,
while [`Iterators.repeated`](@ref)`(iter, n)` is the lazy [`Base.fill`](@ref)`(item, n)`.

!!! compat "Julia 1.11"
    The method `cycle(iter, n)` was added in Julia 1.11.

# Examples
```jldoctest
Repository: JuliaLang/julia
File: base/iterators.jl (Julia)

Code before:
function iterate(ibl::DropWhile)
    y = iterate(ibl.xs)
    while y !== nothing
        ibl.pred(y[1]) || break
        y = iterate(ibl.xs,y[2])
    end
    y
end

IteratorSize(::Type{<:DropWhile}) = SizeUnknown()
eltype(::Type{DropWhile{I,P}}) where {I,P} = eltype(I)
IteratorEltype(::Type{DropWhile{I,P}}) where {I,P} = IteratorEltype(I)

"""
    findeach(f, it)
    findeach(it)

An iterator that generates every key from the key/value pairs of `pairs(it)`,
where `f(value)` returns `true`.

If `f` is not specified, default to `identity`.

`Iterators.findeach` is the lazy equivalent of `findall`.

!!! compat "Julia 1.13"
    Lazy `findeach` requires at least Julia 1.13.

# Examples
```jldoctest
julia> collect(Iterators.findeach(isodd, Dict(2 => 3, 3 => 2)))
1-element Vector{Int64}:
 2

julia> only(Iterators.findeach(==(1), [3,6,2,1]))
4
```
"""
findeach(f, it) = map(first, filter(i -> f(last(i)), pairs(it)))

findeach(it) = findeach(identity, it)


# Cycle an iterator forever

struct Cycle{I}
    xs::I
end

"""
    cycle(iter[, n::Int])


Review comment (suggestion):
The "lazy" qualifier here makes it sound like there's a differently behaving `findeach` in another version.
```suggestion
    `findeach` requires at least Julia 1.13.
```

Code after:
function iterate(ibl::DropWhile)
    y = iterate(ibl.xs)
    while y !== nothing
        ibl.pred(y[1]) || break
        y = iterate(ibl.xs,y[2])
    end
    y
end

IteratorSize(::Type{<:DropWhile}) = SizeUnknown()
eltype(::Type{DropWhile{I,P}}) where {I,P} = eltype(I)
IteratorEltype(::Type{DropWhile{I,P}}) where {I,P} = IteratorEltype(I)

"""
    findeach(f, it)
    findeach(it)

An iterator that generates every key from the key/value pairs of `pairs(it)`,
where `f(value)` returns `true`.

If `f` is not specified, default to `identity`.

`Iterators.findeach` is the lazy equivalent of `findall`.

!!! compat "Julia 1.13"
    `findeach` requires at least Julia 1.13.

# Examples
```jldoctest
julia> collect(Iterators.findeach(isodd, Dict(2 => 3, 3 => 2)))
1-element Vector{Int64}:
 2

julia> only(Iterators.findeach(==(1), [3,6,2,1]))
4
```
"""
findeach(f, it) = (k for (k, v) in pairs(it) if f(v))

findeach(it) = findeach(identity, it)


# Cycle an iterator forever

struct Cycle{I}
    xs::I
end

"""
    cycle(iter[, n::Int])
Repository: JuliaLang/julia
File: stdlib/Random/src/misc.jl (Julia)

Code before:
number generator, see [Random Numbers](@ref).

# Examples
```jldoctest
julia> Random.seed!(3); randstring()
"Lxz5hUwn"

julia> randstring(Xoshiro(3), 'a':'z', 6)
"iyzcsm"

julia> randstring("ACGT")
"TGCTCCTC"
```

!!! note
    `chars` can be any collection of characters, of type `Char` or
    `UInt8` (more efficient), provided [`rand`](@ref) can randomly
    pick characters from it.
"""
function randstring end

let b = UInt8['0':'9';'A':'Z';'a':'z']
    global randstring

    function randstring(r::AbstractRNG, chars=b, n::Integer=8)
        _n = convert(Int, n)
        T = eltype(chars)
        if T === UInt8
            str = Base._string_n(_n)
            GC.@preserve str rand!(r, UnsafeView(pointer(str), _n), chars)
            return str
        else
            v = Vector{T}(undef, _n)
            rand!(r, v, chars)
            return String(v)
        end
    end

    randstring(r::AbstractRNG, n::Integer) = randstring(r, b, n)
    randstring(chars=b, n::Integer=8) = randstring(default_rng(), chars, n)
    randstring(n::Integer) = randstring(default_rng(), b, n)
end


## randsubseq & randsubseq!

# Fill S (resized as needed) with a random subsequence of A, where
# each element of A is included in S with independent probability p.
# (Note that this is different from the problem of finding a random
#  size-m subset of A where m is fixed!)
function randsubseq!(r::AbstractRNG, S::AbstractArray, A::AbstractArray, p::Real)

Review comment (suggestion):
How about this:

```suggestion
        n = convert(Int, n)::Int
```

Then the rest of the diff can be dropped.

Code after:
number generator, see [Random Numbers](@ref).

# Examples
```jldoctest
julia> Random.seed!(3); randstring()
"Lxz5hUwn"

julia> randstring(Xoshiro(3), 'a':'z', 6)
"iyzcsm"

julia> randstring("ACGT")
"TGCTCCTC"
```

!!! note
    `chars` can be any collection of characters, of type `Char` or
    `UInt8` (more efficient), provided [`rand`](@ref) can randomly
    pick characters from it.
"""
function randstring end

let b = UInt8['0':'9';'A':'Z';'a':'z']
    global randstring

    function randstring(r::AbstractRNG, chars, n::Int)
        T = eltype(chars)
        if T === UInt8
            str = Base._string_n(n)
            GC.@preserve str rand!(r, UnsafeView(pointer(str), n), chars)
            return str
        else
            v = Vector{T}(undef, n)
            rand!(r, v, chars)
            return String(v)
        end
    end

    # Zero-arg methods:
    randstring() = randstring(default_rng(), b, 8)
    # One-arg methods:
    randstring(r::AbstractRNG) = randstring(r, b, 8)
    randstring(chars) = randstring(default_rng(), chars, 8)
    randstring(n::Integer) = randstring(default_rng(), b, convert(Int, n))
    # Two-arg methods:
    randstring(r::AbstractRNG, chars) = randstring(r, chars, 8)
    randstring(r::AbstractRNG, n::Integer) = randstring(r, b, convert(Int, n))
    randstring(chars, n::Integer) = randstring(default_rng(), chars, convert(Int, n))
    # Three-arg methods:
    randstring(r::AbstractRNG, chars, n::Integer) = randstring(r, chars, convert(Int, n))
end
Repository: JuliaLang/julia
File: src/support/ios.c (C)

Code before:
        s->bm = bm_none;
    if (fd == STDOUT_FILENO)
        s->bm = bm_line;
    return s;
}

ios_t *ios_stdin = NULL;
ios_t *ios_stdout = NULL;
ios_t *ios_stderr = NULL;
ios_t *ios_safe_stdout = NULL;
ios_t *ios_safe_stderr = NULL;

void ios_init_stdstreams(void)
{
    ios_stdin = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_stdin, STDIN_FILENO, 0, 0);

    ios_stdout = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_stdout, STDOUT_FILENO, 0, 0);
    ios_stdout->bm = bm_line;

    ios_stderr = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_stderr, STDERR_FILENO, 0, 0);
    ios_stderr->bm = bm_none;

    // the following 'safe' variants must use `bm_line` / `bm_none`
    // to avoid memory allocation in an async-signal context

    ios_safe_stdout = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_safe_stdout, STDOUT_FILENO, 0, 0);
    ios_safe_stdout->bm = bm_line;

    ios_safe_stderr = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_safe_stderr, STDERR_FILENO, 0, 0);
    ios_safe_stderr->bm = bm_none;
}

/* higher level interface */

int ios_putc(int c, ios_t *s)
{
    char ch = (char)c;

    if (s->state == bst_wr && s->bpos < s->maxsize && s->bm != bm_none) {
        s->buf[s->bpos++] = ch;
        _write_update_pos(s);
        if (s->bm == bm_line && ch == '\n')
            ios_flush(s);
        return 1;
    }
    return (int)ios_write(s, &ch, 1);

Review comment (suggestion):
`bm_line` seems likely very, very unsafe, so only `bm_none` is okay here. Only when allocated as a local variable, then bm_line is good and works to make something like stderr much faster without sacrificing safety.

Code after:
        s->bm = bm_none;
    if (fd == STDOUT_FILENO)
        s->bm = bm_line;
    return s;
}

ios_t *ios_stdin = NULL;
ios_t *ios_stdout = NULL;
ios_t *ios_stderr = NULL;
ios_t *ios_safe_stderr = NULL;

void ios_init_stdstreams(void)
{
    ios_stdin = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_stdin, STDIN_FILENO, 0, 0);

    ios_stdout = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_stdout, STDOUT_FILENO, 0, 0);
    ios_stdout->bm = bm_line;

    ios_stderr = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_stderr, STDERR_FILENO, 0, 0);
    ios_stderr->bm = bm_none;

    // this 'safe' variant must use `bm_none` to avoid memory allocation
    // in an async-signal context
    ios_safe_stderr = (ios_t*)malloc_s(sizeof(ios_t));
    ios_fd(ios_safe_stderr, STDERR_FILENO, 0, 0);
    ios_safe_stderr->bm = bm_none;
}

/* higher level interface */

int ios_putc(int c, ios_t *s)
{
    char ch = (char)c;

    if (s->state == bst_wr && s->bpos < s->maxsize && s->bm != bm_none) {
        s->buf[s->bpos++] = ch;
        _write_update_pos(s);
        if (s->bm == bm_line && ch == '\n')
            ios_flush(s);
        return 1;
    }
    return (int)ios_write(s, &ch, 1);
}

int ios_getc(ios_t *s)
{
    char ch = 0;
    if (s->state == bst_rd && s->bpos < s->size) {
Repository: JuliaLang/julia
File: base/iterators.jl (Julia)

Code before:
julia> nth(2:2:10, 4)
8

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(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}, n::Integer) = begin
    n > itr.n ? throw(BoundsError(itr, n)) : nth(itr.xs)
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.

Review comment (suggestion):
```suggestion
nth(itr::Take{Repeated{T}}, n::Integer) where T = begin
    n > itr.n ? throw(BoundsError(itr, n)) : nth(itr.xs, n)
end
```

Code after:
```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)
Repository: JuliaLang/julia
File: base/iterators.jl (Julia)

Code before:
julia> nth(2:2:10, 4)
8

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(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.

Review comment (suggestion):
```suggestion
nth(itr::Take{<:Repeated}, n::Integer) =
    n > itr.n ? throw(BoundsError(itr, n)) : nth(itr.xs, n)
```

Code after:
```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)
Repository: JuliaLang/julia
File: base/iterators.jl (Julia)

Code before:
julia> Iterators.nth(2:2:10, 4)
8

julia> Iterators.nth(reshape(1:30, (5,6)), 6)
6

julia> stateful = Iterators.Stateful(1:10); nth(stateful, 7)
7

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.

Review comment (suggestion):
```suggestion
nth(itr::Take{<:Repeated}, n::Integer) =
    n > itr.n ? throw(BoundsError(itr, n)) : nth(itr.xs, n)
```

Code after:
```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)
Repository: JuliaLang/julia
File: base/iterators.jl (Julia)

Code before:
julia> Iterators.nth(2:2:10, 4)
8

julia> Iterators.nth(reshape(1:30, (5,6)), 6)
6

julia> stateful = Iterators.Stateful(1:10); nth(stateful, 7)
7

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.

Review comment (suggestion):
```suggestion
function nth(itr::Cycle{I}, n::Integer) where {I}
```

Code after:
```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)