primebois.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}

{-# OPTIONS_GHC -O2 #-}
{-# OPTIONS_GHC -funbox-strict-fields #-}

{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}

module Prime (primeFast) where

import Gauge.Main
import qualified Data.List as L
import GHC.Int
import GHC.Exts

main :: IO ()
main = do
  let ps = (L.take 50000 (L.cycle [0..1000]))

  defaultMain
    [
      bgroup "prime test"
      [ bench "AKS_SLOW" $ whnf primeMany ps
      , bench "AKS_FAST" $ whnf primeManyFast ps 
      ]
    ]

primeManyFast :: [Int] -> [Bool]
primeManyFast xs = fmap primeFast xs

primeFast :: Int -> Bool
primeFast !p
  | p < 2 = False
  | otherwise = go p 1 1
{-# INLINE primeFast #-}

go :: Int -> Int -> Int -> Bool
go !p !z !i = let !z' = z * (p - i + 1) `unsafeQuot` i in
  if i < p - 1
    then if unsafeRem z' p == 0 || i < 2
      then go p z' (i + 1)
      else False
    else True

unsafeQuot :: Int -> Int -> Int
unsafeRem  :: Int -> Int -> Int
unsafeQuot (I# i) (I# j) = I# (quotInt# i j)
unsafeRem (I# i) (I# j) = I# (remInt# i j)

{-# INLINE unsafeQuot #-}
{-# INLINE unsafeRem #-}

primeMany :: [Int] -> [Bool]
primeMany xs = fmap prime xs

prime :: Int -> Bool
prime !p
  | p < 2     = False
  | otherwise = and [mod n p == 0 | n <- init . tail $ coeffs p]
  where
    coeffs :: Int -> [Int]
    coeffs !p' = scanl (\z i -> z * (p' - i + 1) `div` i) 1 [1..p']