avg

avg.hs

{-# LANGUAGE BangPatterns, ScopedTypeVariables #-}

import Control.Applicative      ( liftA2 )
import Control.DeepSeq          ( deepseq )
import Control.Monad            ( forM_ )
import Control.Monad.ST.Strict  ( runST )
import Criterion                ( bench, bgroup, nf, runAndAnalyse )
import Criterion.Config         ( Config(..), defaultConfig, ljust )
import Criterion.Environment    ( measureEnvironment )
import Criterion.Monad          ( withConfig )
import Data.IORef               ( modifyIORef', newIORef, readIORef )
import Data.List                ( foldl' )
import Data.STRef.Strict        ( modifySTRef', newSTRef, readSTRef )
import Data.Strict.Tuple        ( Pair((:!:)) )
import System.IO.Unsafe         ( unsafePerformIO )
import qualified Control.Foldl  as L

avg1 :: Real a => [a] -> Double
avg1 [] = error "empty"
avg1 (x:xs) = let (s, len :: Int) = loop x 1 xs
              in  realToFrac s / fromIntegral len
  where
    loop s len []     = (s, len)
    loop s len (y:ys) = let !s'   = s + y
                            !len' = len + 1
                        in  loop s' len' ys

avg2 :: Real a => [a] -> Double
avg2 [] = error "empty"
avg2 (x:xs) =
    let (s', len') = runST $ do
            s <- newSTRef x
            len <- newSTRef (1 :: Int)
            forM_ xs $ \y -> do
                modifySTRef' s (+ y)
                modifySTRef' len (+ 1)
            liftA2 (,) (readSTRef s) (readSTRef len)
    in  realToFrac s' / fromIntegral len'

avg3 :: Real a => [a] -> Double
avg3 [] = error "empty"
avg3 (x:xs) =
    let s' :!: len' =
            foldl' (\(s :!: len) y -> s + y :!: len + 1) (x :!: (1 :: Int)) xs
    in  realToFrac s' / fromIntegral len'

avg4 :: Real a => [a] -> Double
avg4 [] = error "empty"
avg4 (x:xs) =
    let (s', len') = unsafePerformIO $ do
            s <- newIORef x
            len <- newIORef (1 :: Int)
            forM_ xs $ \y -> do
                modifyIORef' s (+ y)
                modifyIORef' len (+ 1)
            liftA2 (,) (readIORef s) (readIORef len)
    in  realToFrac s' / fromIntegral len'

avg5 :: Real a => [a] -> Double
avg5 = L.fold (liftA2 d L.sum L.genericLength)
  where d s (len :: Int) = realToFrac s / fromIntegral len

main :: IO ()
main = do
    let n = 10000000 :: Int
        arr = [1 .. n :: Int]
        criterionConfig = defaultConfig { cfgPerformGC  = ljust False
                                        , cfgSamples    = ljust 10
                                        }

    deepseq arr $ return ()

    withConfig criterionConfig $ do
        env <- measureEnvironment
        runAndAnalyse (const True) env $
            bgroup "" [ bench "avg1" (nf avg1 arr)
                      , bench "avg2" (nf avg2 arr)
                      , bench "avg3" (nf avg3 arr)
                      , bench "avg4" (nf avg4 arr)
                      , bench "avg5" (nf avg5 arr)
                      ]