[Haskell-cafe] An unusual prime number generator...

1 Dec 2007

      Hi guys.

I just wrote this prime number generator. It's... unusual. It seems to 
go increadibly fast though. Indeed, I asked it to write the first 
1,000,000 prime numbers to a text file, and scanning the bit array 
seemed to take longer than constructing the bit array! Odd...

Anyway, for your amusement:

module Primes2 (compute_primes) where

import Data.Word
import Data.Array.IO

seive :: IOUArray Word32 Bool -> Word32 -> Word32 -> IO ()
seive grid size p = mapM_ (\n -> writeArray grid n False) [p, 2*p .. size]

next_prime :: IOUArray Word32 Bool -> IO Word32
next_prime grid = work grid 2
  where
    work grid n = do
      p <- readArray grid n
      if p
        then return n
        else work grid (n+1)

copy :: Word32 -> IOUArray Word32 Bool -> Word32 -> IO (IOUArray Word32 
Bool)
copy p grid size = do
  let size' = size * p
  grid' <- newArray (1,size') False

  mapM_
    (\n -> do
      b <- readArray grid n
      if b
        then mapM_ (\x -> writeArray grid' (n + size*x) True) [0..p-1]
        else return ()
    )
    [1..size]

  return grid'

init_grid :: IO (IOUArray Word32 Bool)
init_grid = do
  grid <- newArray (1,6) False
  writeArray grid 1 True
  writeArray grid 5 True
  return grid

compute_primes :: Word32 -> IO [Word32]
compute_primes n = do
    g0 <- init_grid
    ps <- work n g0 6
    return (2:3:ps)
  where
    work top grid size
      | size >= top = do
        p <- next_prime grid
        putStrLn $ "Found   prime: " ++ show p
        if p > top
          then return [p]
          else do
            seive grid size p
            ps <- work top grid size
            return (p:ps)
      | otherwise   = do
        p <- next_prime grid
        putStrLn $ "Seiving prime: " ++ show p
        let size' = p*size
        grid' <- copy p grid size
        seive grid' size' p
        ps <- work top grid' size'
        return (p:ps)

-- Debug...
show_grid :: IOUArray Word32 Bool -> IO ()
show_grid grid = do
  (b0,b1) <- getBounds grid
  mapM_ (\n -> do x <- readArray grid n; if x then putChar '-' else 
putChar '#') [b0..b1]
  putStrLn ":"

[Haskell-cafe] An unusual prime number generator...

Andrew Coppin