Merge pull request #56 from JuliaMath/rf/nextprimes

add nextprimes/prevprimes iterators (fix #55)

Author: mschauer

src/Primes.jl

@@ -4,11 +4,12 @@ module Primes
 
 using Base.Iterators: repeated
 
+import Base: iterate, eltype, IteratorSize, IteratorEltype
 using Base: BitSigned
 using Base.Checked: checked_neg
 
 export isprime, primes, primesmask, factor, ismersenneprime, isrieselprime,
-       nextprime, prevprime, prime, prodfactors, radical, totient
+       nextprime, nextprimes, prevprime, prevprimes, prime, prodfactors, radical, totient
 
 include("factorization.jl")
 
@@ -528,6 +529,11 @@ function totient(n::Integer)
     totient(factor(abs(n)))
 end
 
+# add: checked add (when makes sense), result of same type as first argument
+
+add(n::BigInt, x::Int) = n + x
+add(n::Integer, x::Int) = Base.checked_add(n, oftype(n, x))
+
 # add_! : "may" mutate the Integer argument (only for BigInt currently)
 
 # modify a BigInt in-place
@@ -541,7 +547,7 @@ function add_!(n::BigInt, x::Int)
 end
 
 # checked addition, without mutation
-add_!(n::Integer, x::Int) = Base.checked_add(n, oftype(n, x))
+add_!(n::Integer, x::Int) = add(n, x)
 
 """
     nextprime(n::Integer, i::Integer=1; interval::Integer=1)
@@ -691,4 +697,117 @@ julia> prime(3)
 prime(::Type{T}, i::Integer) where {T<:Integer} = i < 0 ? throw(DomainError(i)) : nextprime(T(2), i)
 prime(i::Integer) = prime(Int, i)
 
+
+struct NextPrimes{T<:Integer}
+    start::T
+end
+
+function iterate(np::NextPrimes, state=np.start)
+    p = nextprime(state)
+    (p, add(p, 1))
+end
+
+IteratorSize(::Type{<:NextPrimes}) = Base.IsInfinite()
+IteratorEltype(::Type{<:NextPrimes}) = Base.HasEltype()
+
+eltype(::Type{NextPrimes{T}}) where {T} = T
+
+"""
+    nextprimes(start::Integer)
+
+Return an iterator over all primes greater than or equal to `start`,
+in ascending order.
+"""
+nextprimes(start::Integer) = NextPrimes(start)
+
+"""
+    nextprimes(T::Type=Int)
+
+Return an iterator over all primes, with type `T`.
+Equivalent to `nextprimes(T(1))`.
+"""
+nextprimes(::Type{T}=Int) where {T<:Integer} = nextprimes(one(T))
+
+"""
+    nextprimes(start::Integer, n::Integer)
+
+Return an array of the first `n` primes greater than or equal to `start`.
+
+# Example
+
+```
+julia> nextprimes(10, 3)
+3-element Array{Int64,1}:
+ 11
+ 13
+ 17
+```
+"""
+nextprimes(start::T, n::Integer) where {T<:Integer} =
+    collect(T, Iterators.take(nextprimes(start), n))
+
+struct PrevPrimes{T<:Integer}
+    start::T
+end
+
+function iterate(np::PrevPrimes, state=np.start)
+    if isone(state)
+        nothing
+    else
+        p = prevprime(state)
+        (p, p-one(p))
+    end
+end
+
+IteratorSize(::Type{<:PrevPrimes}) = Base.SizeUnknown()
+IteratorEltype(::Type{<:PrevPrimes}) = Base.HasEltype()
+
+eltype(::Type{PrevPrimes{T}}) where {T} = T
+
+"""
+    prevprimes(start::Integer)
+
+Return an iterator over all primes less than or equal to `start`,
+in descending order.
+
+# Example
+
+```
+julia> collect(prevprimes(10))
+4-element Array{Int64,1}:
+ 7
+ 5
+ 3
+ 2
+```
+"""
+prevprimes(start::Integer) = PrevPrimes(max(one(start), start))
+
+"""
+    prevprimes(start::Integer, n::Integer)
+
+Return an array of the first `n` primes less than or equal to `start`,
+in descending order. When there are less than `n` primes less than or
+equal to `start`, those primes are returned without an error.
+
+# Example
+
+```
+julia> prevprimes(10, 3)
+3-element Array{Int64,1}:
+ 7
+ 5
+ 3
+
+julia> prevprimes(10, 10)
+3-element Array{Int64,1}:
+ 7
+ 5
+ 3
+ 2
+```
+"""
+prevprimes(start::T, n::Integer) where {T<:Integer} =
+    collect(T, Iterators.take(prevprimes(start), n))
+
 end # module

test/runtests.jl

@@ -405,3 +405,47 @@ for T in (Int, UInt, BigInt)
     @test prodfactors(factor(Set, T(123456))) == 3858
     @test prod(factor(T(123456))) == 123456
 end
+
+@testset "nextprimes(::$T)" for T = (Int32, Int64, BigInt)
+    for (i, p) in enumerate(nextprimes(T))
+        @test nextprime(0, i) == p
+        i > 20 && break
+    end
+    @test nextprimes() == nextprimes(Int)
+    for (i, p) in enumerate(nextprimes(T(5)))
+        @test nextprime(T(5), i) == p
+        i > 20 && break
+    end
+    @test nextprimes(T(5), 10) == [nextprime(T(5), i) for i=1:10]
+    @test nextprimes(1, 1)[1] == nextprimes(2, 1)[1] == 2
+    @test nextprimes(3, 1)[1] == 3
+    @test nextprimes(4, 1)[1] == nextprimes(5, 1)[1] == 5
+    @test eltype(nextprimes(10)) == Int
+    @test eltype(nextprimes(big(10))) == BigInt
+    @test Base.IteratorEltype(nextprimes(10)) == Base.HasEltype()
+    @test Base.IteratorSize(nextprimes(10)) == Base.IsInfinite()
+
+end
+
+
+@testset "prevprimes(::$T)" for T = (Int32, Int64, BigInt)
+    for (i, p) in enumerate(prevprimes(T(500)))
+        @test prevprime(T(500), i) == p
+        i > 20 && break
+    end
+    @test prevprimes(T(500), 10) == [prevprime(T(500), i) for i=1:10]
+    @test prevprimes(6, 1)[1] == prevprimes(5, 1)[1] == 5
+    @test prevprimes(4, 1)[1] == prevprimes(3, 1)[1] == 3
+    @test prevprimes(2, 1)[1] == 2
+    @test isempty(prevprimes(1, 1))
+    let p8 = collect(prevprimes(typemax(Int8)))
+        @test length(p8) == 31
+        @test p8[end] == 2
+        @test p8[1] == 127
+        @test eltype(p8) == Int8
+    end
+    @test eltype(prevprimes(10)) == Int
+    @test eltype(prevprimes(big(10))) == BigInt
+    @test Base.IteratorEltype(prevprimes(10)) == Base.HasEltype()
+    @test Base.IteratorSize(prevprimes(10)) == Base.SizeUnknown()
+end