Ticket #2374: ByteArr.hs

File ByteArr.hs, 1.3 KB (added by dolio, 11 years ago)

Improved byte array benchmark

Line 
1{-# OPTIONS_GHC -fglasgow-exts #-}
2
3module Main (main) where
4
5import Prelude hiding (reverse)
6
7import Control.Monad.ST
8
9import GHC.ST
10import GHC.Base
11import GHC.Prim
12
13import Foreign (sizeOf)
14
15import System.Environment
16
17reverse :: MutableByteArray# s -> Int# -> Int# -> State# s -> State# s
18reverse arr i j s
19  | i <# j    = case readIntArray#  arr i    s of { (# s, ei #) ->
20                case readIntArray#  arr j    s of { (# s, ej #) ->
21                case writeIntArray# arr j ei s of { s ->
22                case writeIntArray# arr i ej s of { s ->
23                reverse arr (i +# 1#) (j -# 1#) s } } } }
24  | otherwise = s
25
26
27
28fill :: MutableByteArray# s -> Int# -> State# s -> State# s
29fill arr n = go 0#
30 where
31 go i s
32   | i <# n    = case writeIntArray# arr i i s of { s ->
33                 go (i +# 1#) s }
34   | otherwise = s
35{-# INLINE fill #-}
36
37bench :: Int -> Int -> ST s ()
38bench (I# k) (I# n) = ST go
39 where
40 go s = case sizeOf (0 :: Int)        of { I# w ->
41        case newByteArray# (n *# w) s of { (# s, arr #) ->
42        case fill arr n s             of { s ->
43        go' arr k s } } }
44 go' arr 0# s = (# s, () #)
45 go' arr k  s = case reverse arr 0# (n -# 1#) s of { s ->
46                go' arr (k -# 1#) s }
47
48
49main = do (k:n:_) <- map read `fmap` getArgs
50          stToIO (bench k n)
51          putStrLn "Done."