Optimise sieve_of_eratosthenes

.spec/math/prime/sieve_of_eratosthenes_spec.lua

@@ -1,12 +1,27 @@
 describe("Sieve of Eratosthenes", function()
 	local sieve = require("math.prime.sieve_of_eratosthenes")
 	local check_primality = require("math.prime.is_prime")
+
+	local function check_sieve(n)
+		local sieve_result = sieve(n)
+		assert.equal(n, #sieve_result)
+		for number, is_prime in ipairs(sieve_result) do
+			assert.equal(check_primality(number), is_prime)
+		end
+	end
+
 	it("works for small numbers", function()
-		for n = 1, 5 do
-			for number, is_prime in ipairs(sieve(n ^ 2 * 1000)) do
-				assert.equal(check_primality(number), is_prime)
-			end
+		for i = 1, 10 do
+			check_sieve(i)
 		end
+		check_sieve(24)
+		check_sieve(25)
+		check_sieve(26)
+		check_sieve(1000)
+		check_sieve(4000)
+		check_sieve(9000)
+		check_sieve(16000)
+		check_sieve(25000)
 	end)
 	it("yields the correct count for large numbers", function()
 		local count = 0
@@ -17,4 +32,9 @@ describe("Sieve of Eratosthenes", function()
 		end
 		assert.equal(78498, count)
 	end)
+	it("should throw error when input is not positive", function()
+		assert.has_error(function()
+			sieve(0)
+		end)
+	end)
 end)

src/math/prime/sieve_of_eratosthenes.lua

@@ -2,13 +2,14 @@
 return function(
 	n -- number
 )
+	assert(n > 0, "n must be positive")
 	local is_prime = { false }
 	for m = 2, n do -- mark as prime
 		is_prime[m] = true
 	end
-	for m = 2, n / 2 do -- iterate possible primes
+	for m = 2, math.sqrt(n) do -- iterate possible primes
 		if is_prime[m] then
-			for l = 2 * m, n, m do -- iterate multiples
+			for l = m * m, n, m do -- iterate multiples
 				is_prime[l] = false -- "cross out" composite
 			end
 		end