sieve.hs

module Main where

import qualified Data.Vector.Unboxed.Mutable as VM
import qualified Data.Vector.Unboxed as V
import Control.Monad.ST (runST)
import Control.Monad (when)

sieve :: Int -> V.Vector Bool
sieve n = runST $ do
  vm <- VM.new n
  VM.set vm True
  let
    go i = when (i < n) $ do
      x <- VM.unsafeRead vm i
      let go' j = when (j < n) (VM.unsafeWrite vm j False *> go' (j + i))
      when x (go' (2 * i))
      go (i + 1)
  go 2
  V.freeze vm

primes :: Int -> [Int]
primes = drop 2 . V.ifoldr (\i x xs -> if x then i : xs else xs) [] . sieve

main :: IO ()
main = print $ primes 1000000