From e24e2f091ad0f9f30ad85bf448130325a3f5c060 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mos=C3=A8=20Giordano?= <giordano@users.noreply.github.com>
Date: Mon, 13 Jul 2020 01:08:24 +0100
Subject: [PATCH] Add two-argument `first` and `last` methods for any iterable
 (#34868)

* Add `{first,last}(::AbstractVector, ::Integer)` methods

* Apply suggestions from code review

Co-Authored-By: Milan Bouchet-Valat <nalimilan@club.fr>

* Add tests for OffsetArrays

* Use `begin` in place of `firstindex`

* Generalise two-argument `first` and `last` to any iterable

Co-authored-by: Milan Bouchet-Valat <nalimilan@club.fr>
---
 NEWS.md               |  2 ++
 base/abstractarray.jl | 54 +++++++++++++++++++++++++++++++++++++++++++
 test/abstractarray.jl | 15 ++++++++++++
 test/offsetarray.jl   | 15 ++++++++++++
 4 files changed, 86 insertions(+)

diff --git a/NEWS.md b/NEWS.md
index d0c2437ed5c0a..dd86f5c8409c1 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -65,6 +65,8 @@ Standard library changes
 * `unique(f, itr; seen=Set{T}())` now allows you to declare the container type used for
   keeping track of values returned by `f` on elements of `itr` ([#36280]).
 * `Libdl` has been moved to `Base.Libc.Libdl`, however it is still accessible as an stdlib ([#35628]).
+* `first` and `last` functions now accept an integer as second argument to get that many
+  leading or trailing elements of any iterable ([#34868]).
 
 #### LinearAlgebra
 * New method `LinearAlgebra.issuccess(::CholeskyPivoted)` for checking whether pivoted Cholesky factorization was successful ([#36002]).
diff --git a/base/abstractarray.jl b/base/abstractarray.jl
index 504b51efec3e5..68dce7e7d0b20 100644
--- a/base/abstractarray.jl
+++ b/base/abstractarray.jl
@@ -343,6 +343,33 @@ function first(itr)
     x[1]
 end
 
+"""
+    first(itr, n::Integer)
+
+Get the first `n` elements of the iterable collection `itr`, or fewer elements if `v` is not
+long enough.
+
+# Examples
+```jldoctest
+julia> first(["foo", "bar", "qux"], 2)
+2-element Vector{String}:
+ "foo"
+ "bar"
+
+julia> first(1:6, 10)
+1:6
+
+julia> first(Bool[], 1)
+Bool[]
+```
+"""
+first(itr, n::Integer) = collect(Iterators.take(itr, n))
+# Faster method for vectors
+function first(v::AbstractVector, n::Integer)
+    n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
+    @inbounds v[begin:min(begin + n - 1, end)]
+end
+
 """
     last(coll)
 
@@ -361,6 +388,33 @@ julia> last([1; 2; 3; 4])
 """
 last(a) = a[end]
 
+"""
+    last(itr, n::Integer)
+
+Get the last `n` elements of the iterable collection `itr`, or fewer elements if `v` is not
+long enough.
+
+# Examples
+```jldoctest
+julia> last(["foo", "bar", "qux"], 2)
+2-element Vector{String}:
+ "bar"
+ "qux"
+
+julia> last(1:6, 10)
+1:6
+
+julia> last(Float64[], 1)
+Float64[]
+```
+"""
+last(itr, n::Integer) = reverse!(collect(Iterators.take(Iterators.reverse(itr), n)))
+# Faster method for arrays
+function last(v::AbstractArray, n::Integer)
+    n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
+    @inbounds v[max(begin, end - n + 1):end]
+end
+
 """
     strides(A)
 
diff --git a/test/abstractarray.jl b/test/abstractarray.jl
index 335697b65f8eb..2db1638c8c950 100644
--- a/test/abstractarray.jl
+++ b/test/abstractarray.jl
@@ -1126,3 +1126,18 @@ end
         end
     end
 end
+
+@testset "first/last n elements of $(typeof(itr))" for itr in (collect(1:9),
+                                                               [1 4 7; 2 5 8; 3 6 9],
+                                                               ntuple(identity, 9))
+    @test first(itr, 6) == [itr[1:6]...]
+    @test first(itr, 25) == [itr[:]...]
+    @test first(itr, 25) !== itr
+    @test first(itr, 1) == [itr[1]]
+    @test_throws ArgumentError first(itr, -6)
+    @test last(itr, 6) == [itr[end-5:end]...]
+    @test last(itr, 25) == [itr[:]...]
+    @test last(itr, 25) !== itr
+    @test last(itr, 1) == [itr[end]]
+    @test_throws ArgumentError last(itr, -6)
+end
diff --git a/test/offsetarray.jl b/test/offsetarray.jl
index d8e2518e2f139..4f8e5224db01b 100644
--- a/test/offsetarray.jl
+++ b/test/offsetarray.jl
@@ -609,3 +609,18 @@ end
     @test_throws DimensionMismatch maximum!(fill(0, -4:-4, 7:7, -6:-5, 1:1), B)
     @test_throws DimensionMismatch minimum!(fill(0, -4:-4, 7:7, -6:-5, 1:1), B)
 end
+
+@testset "first/last n elements of vector" begin
+    v0 = rand(6)
+    v = OffsetArray(v0, (-3,))
+    @test_throws ArgumentError first(v, -2)
+    @test first(v, 2) == v[begin:begin+1]
+    @test first(v, 100) == v0
+    @test first(v, 100) !== v
+    @test first(v, 1) == [v[begin]]
+    @test_throws ArgumentError last(v, -2)
+    @test last(v, 2) == v[end-1:end]
+    @test last(v, 100) == v0
+    @test last(v, 100) !== v
+    @test last(v, 1) == [v[end]]
+end