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