Specialize Iterators functions (#390)

* Specialize Iterators.rest to return an AbstractFill

Specialize Base.rest to return an AbstractFill

* Specialize Iterators.drop and Iterators.take

* Fix error message in take

* Lowercase in error message

* Bump version to v1.14.0

Author: jishnub

Project.toml

@@ -1,6 +1,6 @@
 name = "FillArrays"
 uuid = "1a297f60-69ca-5386-bcde-b61e274b549b"
-version = "1.13.0"
+version = "1.14.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/FillArrays.jl

@@ -229,10 +229,22 @@ Base.@propagate_inbounds getindex(A::AbstractFill, kr::AbstractArray{Bool}) = _f
 
 @inline Base.iterate(F::AbstractFill) = length(F) == 0 ? nothing : (v = getindex_value(F); (v, (v, 1)))
 @inline function Base.iterate(F::AbstractFill, (v, n))
-    n >= length(F) && return nothing
+    1 <= n < length(F) || return nothing
     v, (v, n+1)
 end
 
+# Iterators
+Iterators.rest(F::AbstractFill, (_,n)) = fillsimilar(F, n <= 0 ? 0 : max(length(F)-n, 0))
+function Iterators.drop(F::AbstractFill, n::Integer)
+    n >= 0 || throw(ArgumentError("drop length must be nonnegative"))
+    fillsimilar(F, max(length(F)-n, 0))
+end
+function Iterators.take(F::AbstractFill, n::Integer)
+    n >= 0 || throw(ArgumentError("take length must be nonnegative"))
+    fillsimilar(F, min(n, length(F)))
+end
+Base.rest(F::AbstractFill, s) = Iterators.rest(F, s)
+
 #################
 # Sorting
 #################

test/runtests.jl

@@ -904,6 +904,35 @@ end
     end
 end
 
+@testset "iterators" begin
+    @testset "invalid state" begin
+        @test isnothing(iterate(Ones(4), (1,-3)))
+        @test isempty(Iterators.rest(Ones(4), (1,-3)))
+    end
+    @testset "Iterators.rest" begin
+        @test Iterators.rest(Fill(4, 10), (4, 3)) === Fill(4, 7)
+        # Base.rest
+        a, b... = Fill(3, 4)
+        @test a === 3
+        @test b === Fill(3, 3)
+        a, b... = Ones(3, 4)
+        @test a === 1.0
+        @test b === Ones(11)
+    end
+    @testset "Iterators.drop/take" begin
+        @test Iterators.drop(Fill(4, 10), 3) === Fill(4, 7)
+        @test Iterators.take(Fill(4, 10), 3) === Fill(4, 3)
+        @test Iterators.drop(Fill(4, 10), 0) === Fill(4, 10)
+        @test Iterators.take(Fill(4, 10), 0) === Fill(4, 0)
+        @test Iterators.drop(Fill(4, 10), 11) === Fill(4, 0)
+        @test Iterators.take(Fill(4, 10), 11) === Fill(4, 10)
+        @test_throws ArgumentError Iterators.drop(Fill(4, 10), -11)
+        @test_throws ArgumentError Iterators.take(Fill(4, 10), -11)
+        @test Iterators.drop(Ones(4, 10), 3) === Ones(37)
+        @test Iterators.take(Ones(4, 10), 3) === Ones(3)
+    end
+end
+
 @testset "Broadcast" begin
     x = Fill(5,5)
     @test (.+)(x) ≡ x