Imperative Sieve of Eratosthenes using destructive updates in Haskell

Question

I've implemented the prime sieve of Eratosthenes in Haskell using a mutable array of unboxed Bools instead of the usual pure and immutable datastructures.

My main concern with this code is speed. For large upper bounds, it runs approximately 70% slower than C code implementing the same algorithm. Memory is already optimal (1 bit per number plus some small overhead).

I would also like to know if the code is reasonably idiomatic, considering the fact that is impure, and basically non-functional (meaning imperative). Any nitpicks are appreciated.

import Prelude
import System.Environment (getArgs)
import Data.Array.Unboxed (UArray, (!))
import Data.Array.Storable (newArray, readArray, writeArray, getBounds)
import Data.Array.ST (STUArray, runSTUArray)
import Control.Monad (when)
import Control.Monad.ST (ST)

markMultiplesAsNotPrime :: STUArray s Int Bool -> Int -> ST s ()
markMultiplesAsNotPrime isPrime k = do
  (_, n) <- getBounds isPrime
  kIsPrime <- readArray isPrime k
  when kIsPrime $ do
    let multiples = takeWhile (<=n) $ iterate (+k) (k+k)
    mapM_ (\i -> writeArray isPrime i False) multiples

sieve :: Int -> UArray Int Bool
sieve n | n < 0 = undefined
sieve n = runSTUArray $ do
  isPrime <- newArray (2, n) True :: ST s (STUArray s Int Bool)
  let kMax = isqrt' n
  let ks = takeWhile (<=kMax) . iterate (+1) $ 2
  mapM_ (markMultiplesAsNotPrime isPrime) ks
  return isPrime

isqrt' :: Int -> Int
isqrt' = ceiling . (sqrt :: Double -> Double) . fromIntegral

primesSmallerThanOrEqualTo :: Int -> [Int]
primesSmallerThanOrEqualTo n | n < 2 = []
primesSmallerThanOrEqualTo n = let 
    isPrime = sieve n
  in
    [i | i <- [2 .. n], isPrime ! i]

main :: IO ()
main = do
  args <- getArgs
  mapM_ (print . last . primesSmallerThanOrEqualTo . read) args

In case it is relevant: I compiled and executed the code on an amd64 CPU using "The Glorious Glasgow Haskell Compilation System, version 7.10.3" with the optimization switch "-O2".

Some timings (as measured with the runtime system switch "-s", best of five runs): $$ \begin{array}{4c} n & π(n) & t [s] & \log\left(\frac{t_2}{t_1}\right) \bigg/ \log\left(\frac{π(n_2)}{π(n_1)}\right) \\ 10^4 & 1.23 \cdot 10^3 & 2 \cdot 10^{-3} & - \\ 10^5 & 9.59 \cdot 10^3 & 4 \cdot 10^{-3} & 0.3 \\ 10^6 & 7.85 \cdot 10^4 & 1.0 \cdot 10^{-2} & 0.4 \\ 10^7 & 6.65 \cdot 10^5 & 9.3 \cdot 10^{-2} & 1.04 \\ 10^8 & 5.76 \cdot 10^6 & 1.25 \cdot 10^0 & 1.20 \\ 10^9 & 5.08 \cdot 10^7 & 1.43 \cdot 10^1 & 1.12 \\ 10^{10} & 4.55 \cdot 10^8 & 1.61 \cdot 10^2 & 1.10 \\ 10^{11} & 4.12 \cdot 10^9 & 1.88 \cdot 10^3 & 1.12 \end{array} $$

Answer 1

takeWhile (<=n) $ iterate (+k) (k+k) can be written as [2*k, 3*k..n] which, here, can be replaced with [k*k, k*k+k..n].

Similarly, takeWhile (<= kMax) . iterate (+1) $ 2 can be written as [2..kMax].

Algorithmically, we can work with odds only, and switch to [k*k, k*k+2*k..n] and [3,5..kMax] (special-casing the 2). Taking it further to pre-ignoring the multiples of 3 (5,7...) also, this is known as wheel factorization optimization.

mapM and mapM_, though very idiomatic, can be faster if coded manually as loops (recursive functions). This would also enumerate the numbers manually, instead of creating the interim lists and hoping the compiler won't actually create them.

Using unsafeRead and unsafeWrite should be faster than the safe variants.

Lastly, last . primesSmallerThanOrEqualTo can be replaced with lastPrime,

lastPrime n = let arr = sieve n in
              take 1 [ i | i <- [n, n-1..2], arr ! i]

Even faster is to move the last prime-finding code into the sieve function itself, and return that prime from it, instead of the array.

Answer 2

1. Instead of using takeWhile, you could simply compute largest divisible by k and smaller. That avoids needlessly checking condition for all indices.
2. Instead of using undefined it would be good to provide informative error message with error.
3. STUArray is strict in its elements already. So instead of taking last of primes, you can instead walk the array from the end ([n..1]) and use dropWhile (not . (arr!!) to get the last index that is True.
4. Aligning code blocks would make the code look nicer. (Like vertically aligning =.)