snappy-hs-0.1.2.0: benchmark/Bench.hs

{-# LANGUAGE BangPatterns #-}
module Main where

import Control.DeepSeq (force)
import Control.Exception (evaluate)
import Control.Monad (forM)
import Data.IORef
import qualified Data.ByteString as BS
import Data.Time.Clock.POSIX (getPOSIXTime)
import Data.Word (Word8)
import Snappy
import Text.Printf (printf)

-- | Representative slice of the Google Snappy test corpus plus the original
-- sample. Keep this list in sync with @benchmark/Main.hs@.
corpus :: [FilePath]
corpus =
    [ "./data/Isaac.Newton-Opticks.txt"
    , "./data/html"
    , "./data/alice29.txt"
    , "./data/geo.protodata"
    , "./data/urls.10K"
    , "./data/fireworks.jpeg"
    ]

-- | Iterate each file enough times to process at least this many bytes, so a
-- small file still gets a stable timing signal and a large file does not block.
bytesPerFile :: Int
bytesPerFile = 32 * 1024 * 1024

-- | Distinct inputs per file. A small pool defeats let-floating/CSE of the
-- loop-invariant call while keeping memory bounded for large files.
poolSize :: Int
poolSize = 64

main :: IO ()
main = do
    putStrLn "file                         size      compress        decompress"
    putStrLn "--------------------------------------------------------------------"
    rows <- forM corpus $ \path -> do
        src <- BS.readFile path
        let !n     = BS.length src
            !iters = max poolSize (bytesPerFile `div` max 1 n)
            !mb    = fromIntegral (n * iters) / (1024 * 1024) :: Double
            -- Vary the trailing byte so each pool entry is a distinct value.
            pool   = [ BS.snoc src (fromIntegral k :: Word8) | k <- [0 .. poolSize - 1] ]
        !poolV  <- evaluate (force pool)
        !compV  <- evaluate (force (map compress poolV))
        let inputs   = take iters (cycle poolV)
            compreds = take iters (cycle compV)

        cRef <- newIORef (0 :: Int)
        t0c  <- getPOSIXTime
        mapM_ (\b -> writeIORef cRef =<< evaluate (BS.length (compress b))) inputs
        t1c  <- getPOSIXTime

        dRef <- newIORef (0 :: Int)
        t0d  <- getPOSIXTime
        mapM_ (\c -> case decompress c of
                   Left e   -> error ("decode error in " ++ path ++ ": " ++ show e)
                   Right !d -> writeIORef dRef =<< evaluate (BS.length d)) compreds
        t1d  <- getPOSIXTime

        let compSecs   = realToFrac (t1c - t0c) :: Double
            decompSecs = realToFrac (t1d - t0d) :: Double
        printf "%-26s %8d   %8.1f MB/s   %8.1f MB/s\n"
            (baseName path) n (mb / compSecs) (mb / decompSecs)
        pure (mb, compSecs, decompSecs)

    let totMB   = sum [m | (m, _, _) <- rows]
        totComp = sum [c | (_, c, _) <- rows]
        totDec  = sum [d | (_, _, d) <- rows]
    putStrLn "--------------------------------------------------------------------"
    printf "AGGREGATE                              %8.1f MB/s   %8.1f MB/s\n"
        (totMB / totComp) (totMB / totDec)

baseName :: FilePath -> FilePath
baseName = reverse . takeWhile (/= '/') . reverse