git.iko.soy/streamly
on 42df609 help sign in
streamly/src/Streamly/Examples/MergeSortedStreams.hs
41 lines 1.2 kB history raw
1{-# LANGUAGE FlexibleContexts #-}
3module Streamly.Examples.MergeSortedStreams where
5import Data.Word
6import System.Random (getStdGen, randoms)
7import Data.List (sort)
8import Streamly
9import qualified Streamly.Prelude as A
11getSorted :: MonadIO m => StreamT m Word16
12getSorted = do
13 g <- liftIO getStdGen
14 let ls = take 100000 (randoms g) :: [Word16]
15 foldMapWith (<>) return (sort ls)
17mergeAsync :: (Ord a, MonadAsync m)
18 => StreamT m a -> StreamT m a -> StreamT m a
19mergeAsync a b = do
20 x <- lift $ async a
21 y <- lift $ async b
22 merge x y
24merge :: (Ord a, MonadAsync m) => StreamT m a -> StreamT m a -> StreamT m a
25merge a b = do
26 a1 <- lift $ A.uncons a
27 case a1 of
28 Nothing -> b
29 Just (x, ma) -> do
30 b1 <- lift $ A.uncons b
31 case b1 of
32 Nothing -> return x <> ma
33 Just (y, mb) ->
34 if (y < x)
35 then (return y) <> merge (return x <> ma) mb
36 else (return x) <> merge ma (return y <> mb)
38mergeSortedStreams :: IO ()
39mergeSortedStreams = do
40 xs <- A.toList $ mergeAsync getSorted getSorted
41 putStrLn $ show $ length xs