diff --git a/library/Potoki/Core/Consume.hs b/library/Potoki/Core/Consume.hs --- a/library/Potoki/Core/Consume.hs +++ b/library/Potoki/Core/Consume.hs @@ -348,7 +348,7 @@ Execute a Consume concurrently and consume its results. -} {-# INLINABLE concurrently #-} -concurrently :: Int -> Consume a b -> Consume b c -> Consume a c +concurrently :: NFData b => Int -> Consume a b -> Consume b c -> Consume a c concurrently amount consume1 consume2 = transform (B.concurrently amount (J.consume consume1)) consume2 diff --git a/library/Potoki/Core/Fetch.hs b/library/Potoki/Core/Fetch.hs --- a/library/Potoki/Core/Fetch.hs +++ b/library/Potoki/Core/Fetch.hs @@ -214,11 +214,12 @@ {-# INLINABLE mapFilter #-} mapFilter :: (input -> Maybe output) -> Fetch input -> Fetch output mapFilter mapping (Fetch fetchIO) = + {-# SCC "mapFilter" #-} Fetch $ let loop = do fetch <- fetchIO case mapping <$> fetch of - Just (Just output) -> return (Just output) + Just (Just !output) -> return (Just output) Just Nothing -> loop Nothing -> return Nothing in loop @@ -238,11 +239,12 @@ {-# INLINABLE just #-} just :: Fetch (Maybe element) -> Fetch element just (Fetch fetchIO) = + {-# SCC "just" #-} Fetch $ let loop = do fetch <- fetchIO case fetch of - Just (Just element) -> return (Just element) + Just (Just !element) -> return (Just element) Just (Nothing) -> loop Nothing -> return Nothing in loop @@ -283,7 +285,7 @@ Fetch $ do mbElement <- fetchIO case mbElement of - Just element -> xRay element $> mbElement + Just !element -> xRay element $> mbElement Nothing -> return Nothing {-# INLINE lazyByteStringRef #-} diff --git a/library/Potoki/Core/Prelude.hs b/library/Potoki/Core/Prelude.hs --- a/library/Potoki/Core/Prelude.hs +++ b/library/Potoki/Core/Prelude.hs @@ -115,6 +115,10 @@ ------------------------- import Data.Primitive as Exports +-- deepseq +------------------------- +import Control.DeepSeq as Exports + -------------------------------------------------------------------------------- import qualified Data.Text as A diff --git a/library/Potoki/Core/Produce.hs b/library/Potoki/Core/Produce.hs --- a/library/Potoki/Core/Produce.hs +++ b/library/Potoki/Core/Produce.hs @@ -15,6 +15,7 @@ infiniteMVar, lazyByteString, enumInRange, + mergeOrdering, ) where @@ -247,3 +248,58 @@ Produce $ M.Acquire $ do ref <- newIORef from return (A.enumUntil ref to, return ()) + +{-| +Merge two ordered sequences into one +-} +mergeOrdering :: (a -> a -> Bool) -> Produce a -> Produce a -> Produce a +mergeOrdering compare (Produce produceLeft) (Produce produceRight) = Produce $ do + Fetch fetchLeft <- produceLeft + Fetch fetchRight <- produceRight + leftCache <- liftIO $ newIORef Nothing + rightCache <- liftIO $ newIORef Nothing + return $ Fetch $ do + cachedLeftMaybe <- readIORef leftCache + case cachedLeftMaybe of + Just left -> do + fetchedRightMaybe <- fetchRight + case fetchedRightMaybe of + Just right -> if compare left right + then do + writeIORef leftCache Nothing + writeIORef rightCache (Just right) + return (Just left) + else return (Just right) + Nothing -> do + writeIORef leftCache Nothing + return (Just left) + Nothing -> do + fetchedLeftMaybe <- fetchLeft + case fetchedLeftMaybe of + Just left -> do + cachedRightMaybe <- readIORef rightCache + case cachedRightMaybe of + Just right -> if compare left right + then return (Just left) + else do + writeIORef leftCache (Just left) + writeIORef rightCache Nothing + return (Just right) + Nothing -> do + fetchedRightMaybe <- fetchRight + case fetchedRightMaybe of + Just right -> if compare left right + then do + writeIORef rightCache (Just right) + return (Just left) + else do + writeIORef leftCache (Just left) + return (Just right) + Nothing -> return (Just left) + Nothing -> do + cachedRightMaybe <- readIORef rightCache + case cachedRightMaybe of + Just right -> do + writeIORef rightCache Nothing + return (Just right) + Nothing -> fetchRight diff --git a/library/Potoki/Core/Transform.hs b/library/Potoki/Core/Transform.hs --- a/library/Potoki/Core/Transform.hs +++ b/library/Potoki/Core/Transform.hs @@ -15,7 +15,6 @@ list, vector, batch, - chunk, distinctBy, distinct, executeIO, diff --git a/library/Potoki/Core/Transform/Attoparsec.hs b/library/Potoki/Core/Transform/Attoparsec.hs --- a/library/Potoki/Core/Transform/Attoparsec.hs +++ b/library/Potoki/Core/Transform/Attoparsec.hs @@ -91,7 +91,7 @@ which parses the lines concurrently. -} {-# INLINE parseLineBytesConcurrently #-} -parseLineBytesConcurrently :: Int -> K.Parser a -> Transform ByteString (Either Text a) +parseLineBytesConcurrently :: NFData a => Int -> K.Parser a -> Transform ByteString (Either Text a) parseLineBytesConcurrently concurrency parser = extractLines >>> bufferize concurrency >>> unsafeConcurrently concurrency (arr (mapLeft fromString . K.parseOnly parser)) @@ -101,7 +101,7 @@ which parses the lines concurrently. -} {-# INLINE parseNonEmptyLineBytesConcurrently #-} -parseNonEmptyLineBytesConcurrently :: Int -> K.Parser a -> Transform ByteString (Either Text a) +parseNonEmptyLineBytesConcurrently :: NFData a => Int -> K.Parser a -> Transform ByteString (Either Text a) parseNonEmptyLineBytesConcurrently concurrency parser = extractLinesWithoutTrail >>> filter (not . ByteString.null) >>> bufferize (concurrency * 2) >>> unsafeConcurrently concurrency (arr (mapLeft fromString . K.parseOnly parser)) diff --git a/library/Potoki/Core/Transform/Basic.hs b/library/Potoki/Core/Transform/Basic.hs --- a/library/Potoki/Core/Transform/Basic.hs +++ b/library/Potoki/Core/Transform/Basic.hs @@ -49,26 +49,27 @@ {-# INLINE list #-} list :: Transform [a] a list = - Transform $ \ (A.Fetch fetchListIO) -> M.Acquire $ do + Transform $ \ (A.Fetch fetchListIO) -> liftIO $ do bufferRef <- newIORef [] - return $ (, return ()) $ A.Fetch $ do - buffer <- readIORef bufferRef - case buffer of - headVal : tailVal -> do - writeIORef bufferRef tailVal - return (Just headVal) - _ -> - let - fetchElementIO = do - fetchListIO >>= \case - Nothing -> return Nothing - Just (headVal : tailVal) -> do - writeIORef bufferRef tailVal - return (Just headVal) - _ -> do - writeIORef bufferRef [] - return Nothing - in fetchElementIO + let + fetchFromBufferOr whenEmpty = do + buffer <- readIORef bufferRef + case buffer of + (!head) : tail -> do + writeIORef bufferRef tail + return (Just head) + _ -> whenEmpty + fetchFromSource = do + fetchedList <- fetchListIO + case fetchedList of + Just list -> case list of + (!head) : tail -> do + writeIORef bufferRef tail + return (Just head) + _ -> do + fetchFromSource + Nothing -> return Nothing + in return $ A.Fetch $ fetchFromBufferOr fetchFromSource {-# INLINABLE vector #-} vector :: GenericVector.Vector vector a => Transform (vector a) a @@ -91,13 +92,6 @@ loop Nothing -> return Nothing in loop - -{-| -Alias to "batch". --} -{-# DEPRECATED chunk "Use 'batch' instead" #-} -chunk :: GenericVector.Vector vector a => Int -> Transform a (vector a) -chunk = batch {-| Chunk the stream to vector batches of the given size. diff --git a/library/Potoki/Core/Transform/ByteString.hs b/library/Potoki/Core/Transform/ByteString.hs --- a/library/Potoki/Core/Transform/ByteString.hs +++ b/library/Potoki/Core/Transform/ByteString.hs @@ -26,7 +26,7 @@ -} extractLines :: Transform ByteString ByteString extractLines = - lineList >>> filter (not . null) >>> list + lineList >>> list where lineList = Transform $ \ (A.Fetch fetchIO) -> M.Acquire $ do @@ -60,7 +60,7 @@ extractLinesWithoutTrail :: Transform ByteString ByteString extractLinesWithoutTrail = - lineList >>> filter (not . null) >>> list + lineList >>> list where lineList = Transform $ \ (A.Fetch fetchIO) -> M.Acquire $ do diff --git a/library/Potoki/Core/Transform/Concurrency.hs b/library/Potoki/Core/Transform/Concurrency.hs --- a/library/Potoki/Core/Transform/Concurrency.hs +++ b/library/Potoki/Core/Transform/Concurrency.hs @@ -36,7 +36,7 @@ else return (Just batch) {-# INLINE bufferize #-} -bufferize :: Int -> Transform element element +bufferize :: NFData element => Int -> Transform element element bufferize size = Transform $ \ (A.Fetch fetchIO) -> liftIO $ do buffer <- newTBQueueIO size @@ -46,8 +46,9 @@ loop = do fetchingResult <- fetchIO case fetchingResult of - Just !element -> do - atomically $ writeTBQueue buffer element + Just element -> do + forcedElement <- evaluate (force element) + atomically $ writeTBQueue buffer forcedElement loop Nothing -> atomically $ writeTVar activeVar False in loop @@ -90,7 +91,7 @@ The order of the outputs produced is indiscriminate. -} {-# INLINABLE concurrently #-} -concurrently :: Int -> Transform input output -> Transform input output +concurrently :: NFData output => Int -> Transform input output -> Transform input output concurrently workersAmount transform = if workersAmount == 1 then transform @@ -99,7 +100,7 @@ unsafeConcurrently workersAmount transform {-# INLINE unsafeConcurrently #-} -unsafeConcurrently :: Int -> Transform input output -> Transform input output +unsafeConcurrently :: NFData output => Int -> Transform input output -> Transform input output unsafeConcurrently workersAmount (Transform syncTransformIO) = Transform $ \ fetchIO -> liftIO $ do chan <- newTBQueueIO (workersAmount * 2) @@ -111,8 +112,9 @@ loop = do fetchResult <- fetchIO case fetchResult of - Just !result -> do - atomically (writeTBQueue chan result) + Just result -> do + forcedResult <- evaluate (force result) + atomically (writeTBQueue chan forcedResult) loop Nothing -> atomically (modifyTVar' workersCounter pred) in loop *> finalize @@ -129,7 +131,7 @@ {-| A transform, which fetches the inputs asynchronously on the specified number of threads. -} -async :: Int -> Transform input input +async :: NFData input => Int -> Transform input input async workersAmount = Transform $ \ (A.Fetch fetchIO) -> liftIO $ do chan <- atomically newEmptyTMVar @@ -152,7 +154,7 @@ else return Nothing in atomically (readChan <|> terminate) -concurrentlyWithBatching :: Int -> Int -> Transform a b -> Transform a b +concurrentlyWithBatching :: (NFData a, NFData b) => Int -> Int -> Transform a b -> Transform a b concurrentlyWithBatching batching concurrency transform = batch @Vector batching >>> bufferize concurrency >>> unsafeConcurrently concurrency (vector >>> transform >>> batch @Vector batching) >>> diff --git a/potoki-core.cabal b/potoki-core.cabal --- a/potoki-core.cabal +++ b/potoki-core.cabal @@ -1,5 +1,5 @@ name: potoki-core -version: 2.2.16 +version: 2.3 synopsis: Low-level components of "potoki" description: Provides everything required for building custom instances of @@ -49,6 +49,7 @@ attoparsec >=0.13 && <0.15, base >=4.9 && <5, bytestring ==0.10.*, + deepseq >=1.4.3 && <2, directory >=1.3 && <2, foldl >=1.3 && <2, hashable >=1 && <2, diff --git a/test/Potoki.hs b/test/Potoki.hs --- a/test/Potoki.hs +++ b/test/Potoki.hs @@ -47,6 +47,26 @@ result <- C.produceAndConsume produce (fmap F.length D.concat) assertEqual "" 17400 result , + testCase "mergeOrdering produce testCase" $ do + let + consume = D.list + produceAndConsume list1 list2 = C.produceAndConsume (E.mergeOrdering (\a b -> a <= b) (E.list list1) (E.list list2)) (consume) + assertEqual "" [1,2,3,4,5,6,7,8,9,10] =<< produceAndConsume [1,3,5,7,9] [2,4,6,8,10] + assertEqual "" [1,2,3,4,5,6,7,8,9,10] =<< produceAndConsume [2,4,6,8,10] [1,3,5,7,9] + assertEqual "" [1,2,2,3,3,4,5,6,7,8,8,9,9,10] =<< produceAndConsume [1,2,3,8,9,10] [2,3,4,5,6,7,8,9] + assertEqual "" [1,2,3,4,5,6,7,8,9,10] =<< produceAndConsume [1,2,3] [4,5,6,7,8,9,10] + assertEqual "" [1,2,3,4,5,6,7,8,9,10] =<< produceAndConsume [4,5,6] [1,2,3,7,8,9,10] + assertEqual "" [1,2,3,4,5,6,7,8,9,10] =<< produceAndConsume [1,2,3,7,8,9,10] [4,5,6] + assertEqual "" [1,2,3,4,5,6,7,8,9,10] =<< produceAndConsume [1,2,3,4,5,6,7,8,9,10] [] + assertEqual "" [] =<< produceAndConsume ([] :: [Int]) ([] :: [Int]) + , + testProperty "mergeOrdering produce" $ \ (l1 :: [Int], l2 :: [Int]) -> + let + list1 = sort l1 + list2 = sort l2 + list = sort $ l1 ++ l2 + in list === unsafePerformIO (C.produceAndConsume (E.mergeOrdering (\a b -> a <= b) (E.list list1) (E.list list2)) D.list) + , transformPotoki , parsingPotoki @@ -87,7 +107,7 @@ , testCase "Concurrently" $ do let - list = [1..20000] + list = [1..20000] :: [Int] produce = E.list list transform = A.concurrently 12 $ diff --git a/test/Transform.hs b/test/Transform.hs --- a/test/Transform.hs +++ b/test/Transform.hs @@ -24,11 +24,12 @@ testProperty "Applying chunksOf to list has the same effect as the \"batch\" transform" $ let gen = do list <- listOf (choose (0, 1000 :: Int)) - chunkSize <- frequency [(1000, choose (1, 3)), (100, choose (4, 100)), (1, pure 0)] - return (list, chunkSize) - in forAll gen $ \ (list, chunkSize) -> let - listChunks = if chunkSize < 1 then [] else SplitList.chunksOf chunkSize list - potokiChunks = unsafePerformIO $ C.produceAndTransformAndConsume (E.list list) (rmap toList (A.batch @Vector chunkSize)) D.list + batchSize <- frequency [(1000, choose (1, 3)), (100, choose (4, 100)), (1, pure 0)] + traceShowM (list, batchSize) + return (list, batchSize) + in forAll gen $ \ (list, batchSize) -> let + listChunks = if batchSize < 1 then [] else SplitList.chunksOf batchSize list + potokiChunks = unsafePerformIO $ C.produceAndTransformAndConsume (E.list list) (rmap toList (A.batch @Vector batchSize)) D.list in listChunks === potokiChunks , transformProduce