Ticket #25: Hailstone.hs

module Main where

import Data.Word
import Criterion.Main
import Control.Parallel.Strategies
import Data.Array.Repa.Index
import qualified Data.Array.Repa        as R
import qualified Data.Vector.Unboxed    as V


collatzLen :: Word32 -> Word32 -> Word32
collatzLen c 1 = c
collatzLen c n | n `mod` 2 == 0 = collatzLen (c+1) $ n `div` 2
               | otherwise      = collatzLen (c+1) $ 3*n+1

pmax :: Word32 -> Word32 -> Word32
pmax x n = x `max` collatzLen 1 n

solveAll :: Int -> Int -> Word32
solveAll step n
  = foldl max 1 . parMap r0 solve $ zip steps (tail steps)
  where
    steps            = [2, 2 + fromIntegral step .. fromIntegral n]
    solve (from, to) = foldl pmax 1 [from..to]


-- Solutions
-- ---------

strat :: Int -> Word32
strat n =
  let p = 100   -- chunk size
  in  solveAll (n `div` p) n

repa :: Int -> Word32
repa n = R.foldAll max 1
       . R.map (collatzLen 1)
       . R.fromVector (Z:.n-2)
       $ V.enumFromN 2 (n-2)


-- Main
-- ----

main :: IO ()
main = defaultMain
  [ bgroup "hailstone" [ bench "strat" $ nf strat 1000000
                       , bench "repa"  $ nf repa  1000000 ]
  ]