StaticInt for n-first/last (#289)

* StaticInt for n-first/last

Author: Tokazama

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "6.0.8"
+version = "6.0.9"
 
 [deps]
 ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"

src/indexing.jl

@@ -1,4 +1,27 @@
 
+function known_lastindex(::Type{T}) where {T}
+    if known_offset1(T) === nothing || known_length(T) === nothing
+        return nothing
+    else
+        return known_length(T) - known_offset1(T) + 1
+    end
+end
+known_lastindex(@nospecialize x) = known_lastindex(typeof(x))
+
+@inline static_lastindex(x) = Static.maybe_static(known_lastindex, lastindex, x)
+
+function Base.first(x::AbstractVector, n::StaticInt)
+    @boundscheck n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
+    start = offset1(x)
+    @inbounds x[start:min((start - one(start)) + n, static_lastindex(x))]
+end
+
+function Base.last(x::AbstractVector, n::StaticInt)
+    @boundscheck n < 0 && throw(ArgumentError("Number of elements must be nonnegative"))
+    stop = static_lastindex(x)
+    @inbounds x[max(offset1(x), (stop + one(stop)) - n):stop]
+end
+
 function _is_splat(::Type{I}, i::StaticInt) where {I}
     if dynamic(is_splat_index(field_type(I, i)))
         True()

test/indexing.jl

@@ -192,6 +192,13 @@ end
     end
 end
 
+@testset "n-first/last" begin
+    x = MArray([1, 2, 3, 4])
+    n = static(2)
+    @test @inferred(first(x, n)) == [1, 2]
+    @test @inferred(last(x, n)) == [3, 4]
+end
+
 A = zeros(3, 4, 5);
 A[:] = 1:60
 Ap = @view(PermutedDimsArray(A, (3, 1, 2))[:, 1:2, 1])';
@@ -222,4 +229,3 @@ izip = zip(S, S)
 
 sv5 = MArray(zeros(5));
 v5 = Vector{Float64}(undef, 5);
-