carray-0.1.5.2: tests/nsieve-bits-u.hs

{-# OPTIONS -O2 -optc-O #-}
{-# LANGUAGE BangPatterns#-}

--
-- The Computer Language Shootout
-- http://shootout.alioth.debian.org/
--
-- Contributed by Don Stewart
-- nsieve over an ST monad Bool array
--

import Control.Monad.ST
import Data.Array.ST
import Data.Array.Base
import System.Environment
import Control.Monad
import Data.Bits
import Text.Printf

main = do
    n <- getArgs >>= readIO . head :: IO Int
    mapM_ (sieve . (10000 *) . (2 ^)) [n, n-1, n-2]

sieve n = do
   let r = runST (do a <- newArray (2,n) True :: ST s (STUArray s Int Bool)
                     go a n 2 0)
   printf "Primes up to %8d %8d\n" (n::Int) (r::Int) :: IO ()

go !a !m !n !c
    | n == m    = return c
    | otherwise = do
          e <- unsafeRead a n
          if e then let loop !j
                          | j < m     = do
                              x <- unsafeRead a j
                              when x $ unsafeWrite a j False
                              loop (j+n)

                          | otherwise = go a m (n+1) (c+1)
                    in loop (n `shiftL` 1)
               else go a m (n+1) c