git.iko.soy/streamly
on 33ae318 help sign in
streamly/src/Streamly/Examples/MergeSortedStreams.hs
42 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 Streamly.Prelude (once)
10import qualified Streamly.Prelude as A
12getSorted :: Stream Word16
13getSorted = do
14 g <- once getStdGen
15 let ls = take 100000 (randoms g) :: [Word16]
16 foldMap return (sort ls)
18-- | merge two streams generating the elements from each in parallel
19mergeAsync :: Ord a => Stream a -> Stream a -> Stream a
20mergeAsync a b = do
21 x <- once $ async a
22 y <- once $ async b
23 merge x y
25merge :: Ord a => Stream a -> Stream a -> Stream a
26merge a b = do
27 a1 <- once $ A.uncons a
28 case a1 of
29 Nothing -> b
30 Just (x, ma) -> do
31 b1 <- once $ A.uncons b
32 case b1 of
33 Nothing -> return x <> ma
34 Just (y, mb) ->
35 if (y < x)
36 then (return y) <> merge (return x <> ma) mb
37 else (return x) <> merge ma (return y <> mb)
39mergeSortedStreams :: IO ()
40mergeSortedStreams = do
41 xs <- A.toList $ mergeAsync getSorted getSorted
42 putStrLn $ show $ length xs