Fold integer-simple.git into ghc.git (re #8545)
[ghc.git] / testsuite / tests / perf / should_run / T5113.hs
1 {-# LANGUAGE BangPatterns #-}
2 module Main where
3
4 import Data.Array.Base (unsafeRead, unsafeWrite)
5 import Data.Array.ST
6 import Data.Array.Unboxed
7 import Control.Monad.ST
8
9 main = print (divisorCounts 1000000 ! 342)
10
11 isqrt :: Int -> Int
12 isqrt n = floor (sqrt $ fromIntegral n)
13
14 divisorCounts :: Int -> UArray Int Int
15 divisorCounts n = runSTUArray $ do
16 let !rt = isqrt n
17 darr <- newArray (0,n) 1 :: ST s (STUArray s Int Int)
18 let inc i = unsafeRead darr i >>= \k -> unsafeWrite darr i (k+1)
19 note step i
20 | i > n = return ()
21 | otherwise = do
22 inc i
23 note step (i+step)
24 count j
25 | j > rt = return ()
26 | otherwise = do
27 note (2*j) (j*j)
28 count (j+2)
29 note 2 4
30 count 3
31 return darr