primes.lua

View source code here on GitHub!

primes(stop)
Returns:

An iterator over the prime numbers, optionally up to stop

Return type:

table

prime_factors(n)
Returns:

An iterator over the prime factors of n

Return type:

table

 1local function primes(stop)
 2    local prime = 1
 3    local candidate = 2
 4    local sieve = {}
 5
 6    local function next_prime(cand)
 7        local steps = sieve[cand]
 8
 9        if steps then
10            for _, step in ipairs(steps) do
11                local value = cand + step
12                if sieve[value] then
13                    table.insert(sieve[value], step)
14                else
15                    sieve[value] = { step }
16                end
17            end
18            sieve[cand] = nil
19            return next_prime(cand + 1)
20        else
21            sieve[cand * cand] = { cand }
22            return cand
23        end
24    end
25
26    local function next()
27        if stop and prime >= stop then
28            sieve = {}
29            return nil, "Tried to exceed given limit"
30        end
31
32        prime = next_prime(candidate)
33        candidate = prime + 1
34
35        if stop and prime >= stop then
36            sieve = {}
37            return nil, "Tried to exceed given limit"
38        end
39
40        return prime
41    end
42
43    return {
44        next = next
45    }
46end
47
48local function prime_factors(n)
49    local pgen = primes()
50    local p = pgen.next()
51
52    local function next()
53        if n == 1 then
54            return nil
55        end
56
57        while p and n % p ~= 0 do
58            p = pgen.next()
59        end
60
61        if p == nil then
62            return nil
63        end
64
65        if n < 0 then
66            n = -n
67            return -1
68        end
69
70        n = math.floor(n / p)
71        return p
72    end
73
74    return {
75        next = next,
76    }
77end
78
79return {
80    primes = primes,
81    prime_factors = prime_factors,
82}