1{-# LANGUAGE BangPatterns #-} 2{-# LANGUAGE ExistentialQuantification #-} 3{-# LANGUAGE FlexibleContexts #-} 4{-# LANGUAGE MultiParamTypeClasses #-} 5{-# LANGUAGE TupleSections #-} 6{-# LANGUAGE RankNTypes #-} 7-- Collection of three benchmarks: a simple integral sum, monte carlo analysis, 8-- and sliding vector. 9import Control.DeepSeq 10import Control.Monad (foldM) 11import Control.Monad (when, liftM) 12import Control.Monad.IO.Class (liftIO) 13import Gauge.Main 14import Data.Conduit 15import qualified Data.Conduit.Internal as CI 16import qualified Data.Conduit.List as CL 17import qualified Data.Foldable as F 18import Data.IORef 19import Data.List (foldl') 20import Data.Monoid (mempty) 21import qualified Data.Sequence as Seq 22import qualified Data.Vector as VB 23import qualified Data.Vector.Generic as V 24import qualified Data.Vector.Generic.Mutable as VM 25import qualified Data.Vector.Unboxed as VU 26import System.Environment (withArgs) 27import qualified System.Random.MWC as MWC 28import Test.Hspec 29 30data TestBench = TBGroup String [TestBench] 31 | TBBench Benchmark 32 | forall a b. (Eq b, Show b) => TBPure String a b (a -> b) 33 | forall a. (Eq a, Show a) => TBIO String a (IO a) 34 | forall a. (Eq a, Show a) => TBIOTest String (a -> IO ()) (IO a) 35 | forall a. (Eq a, Show a) => TBIOBench String a (IO a) (IO ()) 36 37toSpec :: TestBench -> Spec 38toSpec (TBGroup name tbs) = describe name $ mapM_ toSpec tbs 39toSpec (TBBench _) = return () 40toSpec (TBPure name a b f) = it name $ f a `shouldBe` b 41toSpec (TBIO name a f) = it name $ f >>= (`shouldBe` a) 42toSpec (TBIOTest name spec f) = it name $ f >>= spec 43toSpec (TBIOBench name a f _) = it name $ f >>= (`shouldBe` a) 44 45toBench :: TestBench -> Benchmark 46toBench (TBGroup name tbs) = bgroup name $ map toBench tbs 47toBench (TBBench b) = b 48toBench (TBPure name a _ f) = bench name $ whnf f a 49toBench (TBIO name _ f) = bench name $ whnfIO f 50toBench (TBIOTest name _ f) = bench name $ whnfIO f 51toBench (TBIOBench name _ _ f) = bench name $ whnfIO f 52 53runTestBench :: [TestBench] -> IO () 54runTestBench tbs = do 55 withArgs [] $ hspec $ mapM_ toSpec tbs 56 defaultMain $ map toBench tbs 57 58main :: IO () 59main = runTestBench =<< sequence 60 [ sumTB 61 , mapSumTB 62 , monteCarloTB 63 , fmap (TBGroup "sliding window") $ sequence 64 [ slidingWindow 10 65 , slidingWindow 30 66 , slidingWindow 100 67 , slidingWindow 1000 68 ] 69 ] 70 71----------------------------------------------------------------------- 72 73sumTB :: IO TestBench 74sumTB = do 75 upperRef <- newIORef upper0 76 return $ TBGroup "sum" 77 [ TBPure "Data.List.foldl'" upper0 expected 78 $ \upper -> foldl' (+) 0 [1..upper] 79 , TBIO "Control.Monad.foldM" expected $ do 80 upper <- readIORef upperRef 81 foldM plusM 0 [1..upper] 82 , TBPure "low level" upper0 expected $ \upper -> 83 let go x !t 84 | x > upper = t 85 | otherwise = go (x + 1) (t + x) 86 in go 1 0 87 , TBIO "boxed vectors, I/O" expected $ do 88 upper <- readIORef upperRef 89 VB.foldM' plusM 0 $ VB.enumFromTo 1 upper 90 , TBPure "boxed vectors" upper0 expected 91 $ \upper -> VB.foldl' (+) 0 (VB.enumFromTo 1 upper) 92 , TBPure "unboxed vectors" upper0 expected 93 $ \upper -> VU.foldl' (+) 0 (VU.enumFromTo 1 upper) 94 , TBPure "conduit, pure, fold" upper0 expected 95 $ \upper -> runConduitPure $ CL.enumFromTo 1 upper .| CL.fold (+) 0 96 , TBPure "conduit, pure, foldM" upper0 expected 97 $ \upper -> runConduitPure $ CL.enumFromTo 1 upper .| CL.foldM plusM 0 98 , TBIO "conduit, IO, fold" expected $ do 99 upper <- readIORef upperRef 100 runConduit $ CL.enumFromTo 1 upper .| CL.fold (+) 0 101 , TBIO "conduit, IO, foldM" expected $ do 102 upper <- readIORef upperRef 103 runConduit $ CL.enumFromTo 1 upper .| CL.foldM plusM 0 104 ] 105 where 106 upper0 = 10000 :: Int 107 expected = sum [1..upper0] 108 109 plusM x y = return $! x + y 110 111----------------------------------------------------------------------- 112 113mapSumTB :: IO TestBench 114mapSumTB = return $ TBGroup "map + sum" 115 [ TBPure "boxed vectors" upper0 expected 116 $ \upper -> VB.foldl' (+) 0 117 $ VB.map (+ 1) 118 $ VB.map (* 2) 119 $ VB.enumFromTo 1 upper 120 , TBPure "unboxed vectors" upper0 expected 121 $ \upper -> VU.foldl' (+) 0 122 $ VU.map (+ 1) 123 $ VU.map (* 2) 124 $ VU.enumFromTo 1 upper 125 , TBPure "conduit, connect1" upper0 expected $ \upper -> runConduitPure 126 $ CL.enumFromTo 1 upper 127 .| CL.map (* 2) 128 .| CL.map (+ 1) 129 .| CL.fold (+) 0 130 ] 131 where 132 upper0 = 10000 :: Int 133 expected = sum $ map (+ 1) $ map (* 2) [1..upper0] 134 135----------------------------------------------------------------------- 136 137monteCarloTB :: IO TestBench 138monteCarloTB = return $ TBGroup "monte carlo" 139 [ TBIOTest "conduit" closeEnough $ do 140 gen <- MWC.createSystemRandom 141 successes <- runConduit 142 $ CL.replicateM count (MWC.uniform gen) 143 .| CL.fold (\t (x, y) -> 144 if (x*x + y*(y :: Double) < 1) 145 then t + 1 146 else t) 147 (0 :: Int) 148 return $ fromIntegral successes / fromIntegral count * 4 149 , TBIOTest "low level" closeEnough $ do 150 gen <- MWC.createSystemRandom 151 let go :: Int -> Int -> IO Double 152 go 0 !t = return $! fromIntegral t / fromIntegral count * 4 153 go i !t = do 154 (x, y) <- MWC.uniform gen 155 let t' 156 | x*x + y*(y :: Double) < 1 = t + 1 157 | otherwise = t 158 go (i - 1) t' 159 go count (0 :: Int) 160 ] 161 where 162 count = 100000 :: Int 163 164 closeEnough x 165 | abs (x - 3.14159 :: Double) < 0.2 = return () 166 | otherwise = error $ "Monte carlo analysis too inaccurate: " ++ show x 167 168----------------------------------------------------------------------- 169 170slidingWindow :: Int -> IO TestBench 171slidingWindow window = do 172 upperRef <- newIORef upper0 173 return $ TBGroup (show window) 174 [ TBIOBench "low level, Seq" expected 175 (swLowLevelSeq window upperRef id (\x y -> x . (F.toList y:)) ($ [])) 176 (swLowLevelSeq window upperRef () (\() y -> rnf y) id) 177 , TBIOBench "conduit, Seq" expected 178 (swConduitSeq window upperRef id (\x y -> x . (F.toList y:)) ($ [])) 179 (swConduitSeq window upperRef () (\() y -> rnf y) id) 180 {- https://ghc.haskell.org/trac/ghc/ticket/9446 181 , TBIOBench "low level, boxed Vector" expected 182 (swLowLevelVector window upperRef id (\x y -> x . (VB.toList y:)) ($ [])) 183 (swLowLevelVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id) 184 -} 185 , TBBench $ bench "low level, boxed Vector" $ whnfIO $ 186 swLowLevelVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id 187 188 {- https://ghc.haskell.org/trac/ghc/ticket/9446 189 , TBIOBench "conduit, boxed Vector" expected 190 (swConduitVector window upperRef id (\x y -> x . (VB.toList y:)) ($ [])) 191 (swConduitVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id) 192 -} 193 194 , TBBench $ bench "conduit, boxed Vector" $ whnfIO $ 195 swConduitVector window upperRef () (\() y -> rnf (y :: VB.Vector Int)) id 196 197 198 , TBIOBench "low level, unboxed Vector" expected 199 (swLowLevelVector window upperRef id (\x y -> x . (VU.toList y:)) ($ [])) 200 (swLowLevelVector window upperRef () (\() y -> rnf (y :: VU.Vector Int)) id) 201 , TBIOBench "conduit, unboxed Vector" expected 202 (swConduitVector window upperRef id (\x y -> x . (VU.toList y:)) ($ [])) 203 (swConduitVector window upperRef () (\() y -> rnf (y :: VU.Vector Int)) id) 204 ] 205 where 206 upper0 = 10000 207 expected = 208 loop [1..upper0] 209 where 210 loop input 211 | length x == window = x : loop y 212 | otherwise = [] 213 where 214 x = take window input 215 y = drop 1 input 216 217swLowLevelSeq :: Int -> IORef Int -> t -> (t -> Seq.Seq Int -> t) -> (t -> t') -> IO t' 218swLowLevelSeq window upperRef t0 f final = do 219 upper <- readIORef upperRef 220 221 let phase1 i !s 222 | i > window = phase2 i s t0 223 | otherwise = phase1 (i + 1) (s Seq.|> i) 224 225 phase2 i !s !t 226 | i > upper = t' 227 | otherwise = phase2 (i + 1) s' t' 228 where 229 t' = f t s 230 s' = Seq.drop 1 s Seq.|> i 231 232 return $! final $! phase1 1 mempty 233 234swLowLevelVector :: V.Vector v Int 235 => Int 236 -> IORef Int 237 -> t 238 -> (t -> v Int -> t) 239 -> (t -> t') 240 -> IO t' 241swLowLevelVector window upperRef t0 f final = do 242 upper <- readIORef upperRef 243 244 let go !i !t _ _ _ | i > upper = return $! final $! t 245 go !i !t !end _mv mv2 | end == bufSz = newBuf >>= go i t sz mv2 246 go !i !t !end mv mv2 = do 247 VM.unsafeWrite mv end i 248 when (end > sz) $ VM.unsafeWrite mv2 (end - sz) i 249 let end' = end + 1 250 t' <- 251 if end' < sz 252 then return t 253 else do 254 v <- V.unsafeFreeze $ VM.unsafeSlice (end' - sz) sz mv 255 return $! f t v 256 go (i + 1) t' end' mv mv2 257 258 mv <- newBuf 259 mv2 <- newBuf 260 go 1 t0 0 mv mv2 261 where 262 sz = window 263 bufSz = 2 * window 264 newBuf = VM.new bufSz 265 266swConduitSeq :: Int 267 -> IORef Int 268 -> t 269 -> (t -> Seq.Seq Int -> t) 270 -> (t -> t') 271 -> IO t' 272swConduitSeq window upperRef t0 f final = do 273 upper <- readIORef upperRef 274 275 t <- runConduit 276 $ CL.enumFromTo 1 upper 277 .| slidingWindowC window 278 .| CL.fold f t0 279 return $! final t 280 281swConduitVector :: V.Vector v Int 282 => Int 283 -> IORef Int 284 -> t 285 -> (t -> v Int -> t) 286 -> (t -> t') 287 -> IO t' 288swConduitVector window upperRef t0 f final = do 289 upper <- readIORef upperRef 290 291 t <- runConduit 292 $ CL.enumFromTo 1 upper 293 .| slidingVectorC window 294 .| CL.fold f t0 295 return $! final t 296 297slidingWindowC :: Monad m => Int -> ConduitT a (Seq.Seq a) m () 298slidingWindowC = slidingWindowCC 299{-# INLINE [0] slidingWindowC #-} 300{-# RULES "unstream slidingWindowC" 301 forall i. slidingWindowC i = CI.unstream (CI.streamConduit (slidingWindowCC i) (slidingWindowS i)) 302 #-} 303 304slidingWindowCC :: Monad m => Int -> ConduitT a (Seq.Seq a) m () 305slidingWindowCC sz = 306 go sz mempty 307 where 308 goContinue st = await >>= 309 maybe (return ()) 310 (\x -> do 311 let st' = st Seq.|> x 312 yield st' >> goContinue (Seq.drop 1 st') 313 ) 314 go 0 st = yield st >> goContinue (Seq.drop 1 st) 315 go !n st = CL.head >>= \m -> 316 case m of 317 Nothing | n < sz -> yield st 318 | otherwise -> return () 319 Just x -> go (n-1) (st Seq.|> x) 320{-# INLINE slidingWindowCC #-} 321 322slidingWindowS :: Monad m => Int -> CI.Stream m a () -> CI.Stream m (Seq.Seq a) () 323slidingWindowS sz (CI.Stream step ms0) = 324 CI.Stream step' $ liftM (\s -> Left (s, sz, mempty)) ms0 325 where 326 step' (Left (s, 0, st)) = return $ CI.Emit (Right (s, st)) st 327 step' (Left (s, i, st)) = do 328 res <- step s 329 return $ case res of 330 CI.Stop () -> CI.Stop () 331 CI.Skip s' -> CI.Skip $ Left (s', i, st) 332 CI.Emit s' a -> CI.Skip $ Left (s', i - 1, st Seq.|> a) 333 step' (Right (s, st)) = do 334 res <- step s 335 return $ case res of 336 CI.Stop () -> CI.Stop () 337 CI.Skip s' -> CI.Skip $ Right (s', st) 338 CI.Emit s' a -> 339 let st' = Seq.drop 1 st Seq.|> a 340 in CI.Emit (Right (s', st')) st' 341{-# INLINE slidingWindowS #-} 342 343slidingVectorC :: V.Vector v a => Int -> ConduitT a (v a) IO () 344slidingVectorC = slidingVectorCC 345{-# INLINE [0] slidingVectorC #-} 346{-# RULES "unstream slidingVectorC" 347 forall i. slidingVectorC i = CI.unstream (CI.streamConduit (slidingVectorCC i) (slidingVectorS i)) 348 #-} 349 350slidingVectorCC :: V.Vector v a => Int -> ConduitT a (v a) IO () 351slidingVectorCC sz = do 352 mv <- newBuf 353 mv2 <- newBuf 354 go 0 mv mv2 355 where 356 bufSz = 2 * sz 357 newBuf = liftIO (VM.new bufSz) 358 359 go !end _mv mv2 | end == bufSz = newBuf >>= go sz mv2 360 go !end mv mv2 = do 361 mx <- await 362 case mx of 363 Nothing -> when (end > 0 && end < sz) $ do 364 v <- liftIO $ V.unsafeFreeze $ VM.take end mv 365 yield v 366 Just x -> do 367 liftIO $ do 368 VM.unsafeWrite mv end x 369 when (end > sz) $ VM.unsafeWrite mv2 (end - sz) x 370 let end' = end + 1 371 when (end' >= sz) $ do 372 v <- liftIO $ V.unsafeFreeze $ VM.unsafeSlice (end' - sz) sz mv 373 yield v 374 go end' mv mv2 375 376slidingVectorS :: V.Vector v a => Int -> CI.Stream IO a () -> CI.Stream IO (v a) () 377slidingVectorS sz (CI.Stream step ms0) = 378 CI.Stream step' ms1 379 where 380 bufSz = 2 * sz 381 newBuf = liftIO (VM.new bufSz) 382 383 ms1 = do 384 s <- ms0 385 mv <- newBuf 386 mv2 <- newBuf 387 return (s, 0, mv, mv2) 388 389 step' (_, -1, _, _) = return $ CI.Stop () 390 step' (s, end, _mv, mv2) | end == bufSz = do 391 mv3 <- newBuf 392 return $ CI.Skip (s, sz, mv2, mv3) 393 step' (s, end, mv, mv2) = do 394 res <- step s 395 case res of 396 CI.Stop () 397 | end > 0 && end < sz -> do 398 v <- liftIO $ V.unsafeFreeze $ VM.take end mv 399 return $ CI.Emit (s, -1, mv, mv2) v 400 | otherwise -> return $ CI.Stop () 401 CI.Skip s' -> return $ CI.Skip (s', end, mv, mv2) 402 CI.Emit s' x -> liftIO $ do 403 VM.unsafeWrite mv end x 404 when (end > sz) $ VM.unsafeWrite mv2 (end - sz) x 405 let end' = end + 1 406 state = (s', end', mv, mv2) 407 if end' >= sz 408 then do 409 v <- V.unsafeFreeze $ VM.unsafeSlice (end' - sz) sz mv 410 return $ CI.Emit state v 411 else return $ CI.Skip state 412{-# INLINE slidingVectorS #-} 413