Repa histogram benchmark

histogram.hs

{-# LANGUAGE FlexibleInstances #-}
import Codec.Picture.Repa
import qualified Data.Array.Repa as R
import qualified Data.Array.Repa.Unsafe as RU
import qualified Data.Vector.Unboxed as VU
import Data.Array.Repa.Repr.ForeignPtr (F)
import Data.Array.Repa ((:.), Array, (:.)(..), Z(..), DIM3)
import System.Environment
import Data.Word
import Criterion.Main


main:: IO ()
main = do
    path: xs <- getArgs
    eImg <- readImageRGBA path
    case eImg of
      Left e  -> error e
      Right i -> withArgs xs $ defaultMain [bench "JuicyPixels-repa histogram" $ whnf (histograms.imgData) i]


type Histogram = R.Array R.D R.DIM1 Int
type HistogramRGBA = (Histogram, Histogram, Histogram, Histogram)

instance Benchmarkable HistogramRGBA where
    run t@(r,g,b,a) n 
        | n <= 0    = return ()
        | otherwise = return (map R.computeUnboxedS  [r,g,b,a]) >> run t (n-1)
    {-# INLINE run #-}

histograms :: Array F DIM3 Word8 -> (Histogram, Histogram, Histogram, Histogram)
histograms arr =
    let (Z:.nrRow:.nrCol:._) = R.extent arr
        zero = R.fromFunction (Z:.256) (\_ -> 0 :: Int)
        incElem idx x = RU.unsafeTraverse x id (\l i -> l i + if i==(Z:.fromIntegral idx) then 1 else 0)
    in Prelude.foldl (\(hR, hG, hB, hA) (row,col) ->
             let r = R.unsafeIndex arr (Z:.row:.col:.0)
                 g = R.unsafeIndex arr (Z:.row:.col:.1)
                 b = R.unsafeIndex arr (Z:.row:.col:.2)
                 a = R.unsafeIndex arr (Z:.row:.col:.3)
             in (incElem r hR, incElem g hG, incElem b hB, incElem a hA))
          (zero,zero,zero,zero)
          [ (row,col) | row <- [0..nrRow-1], col <- [0..nrCol-1] ]