Prime sieve

Sieve.hs

{-
A direct implementation of a prime sieve ...
-}

module ArrSieve
  ( Sieve(unSieve)
  , isPrime
  , unsafeIsPrime
  , sieve
  ) where

import Data.Array.MArray
import Data.Array.ST
import Data.Array.Unboxed
import Data.Word

------------------------------------------------------------------------

{-|
A prime sieve, where indices map to a flag indicating whether
the corresponding number is composite or prime.
-}
newtype Sieve = Sieve { unSieve :: UArray Word64 Bool }

{-|
@
isPrime s n
@
tests @n@ for primeness with bounds checking, returning
@Nothing@ if @n@ is out of bounds.
-}
isPrime :: Sieve -> Integer -> Maybe Bool
isPrime (Sieve v) = f . fromIntegral
  where f i = if inRange (bounds v) i then Just (not $ v ! i) else Nothing

{-|
@
unsafeIsPrime s n
@
tests @n@ for primeness without bounds checking.

Typical use-case

@
foo ... = do
  let f = unsafeIsPrime (sieve someLargeNumber)
  ...
  forM_ [2..someLargeNumber] $ \n -> do
    putStrLn (show n ++ " " ++ if f n then "is" else "isn't" ++ " prime")
@
-}
unsafeIsPrime :: Sieve -> Integer -> Bool
unsafeIsPrime (Sieve v) = not . (v !) . fromIntegral

{-|
@
sieve n
@
returns a sieve for prime numbers up to @n@.

The computation errors if @n > (maxBound :: Word64)@.
-}
sieve :: Integer -> Sieve
sieve n | n > fromIntegral (maxBound :: Word64) = error "sieve: n is too large"
sieve n = Sieve $ runSTUArray $ do
  let n' = fromIntegral n
  v <- newArray (2, n') False
  let loop i = do
        mapM_ (flip (writeArray v) True) (drop 1 $ enumFromThenTo i (i + i) n')
        maybe (return ()) loop =<< firstM (fmap not . readArray v) [i + 1 .. n']
  loop 2
  return v

{-|
@
firstM (const $ return True) (1 : repeat undefined) = return $ Just 1
@
-}
firstM :: (Monad m) => (a -> m Bool) -> [a] -> m (Maybe a)
firstM f = go where
  go []     = return Nothing
  go (x:xs) = do
    b <- f x
    if b then return (Just x) else go xs