@@ -216,15 +216,24 @@ module Streamly.Internal.Data.MutArray.Type
216216 , pinnedChunksOf
217217 , buildChunks
218218
219− -- *** Flatten a stream of arrays
220− , concatChunks
221− , concatChunksRev
222−
223219 -- *** Split an array into slices
224220 -- , getSlicesFromLenN
225221 , splitOn
226222 -- , slicesOf
227223
224+ -- *** Flatten a stream of arrays
225+ , concatChunks
226+ , concatChunksRev
227+
228+ -- *** Compaction
229+ , SpliceState (..)
230+ , pCompactChunksLE
231+ , rCompactChunksLE
232+ , fCompactChunksGE
233+ , lCompactChunksGE
234+ , compactChunksGE
235+ , compactChunksEQ
236+
228237 -- ** Utilities
229238 , roundUpToPower2
230239 , memcpy
@@ -248,6 +257,7 @@ where
248257
249258import Control.Monad (when, void)
250259import Control.Monad.IO.Class (MonadIO(..))
260+import Data.Bifunctor (first)
251261import Data.Bits (shiftR, (.|.), (.&.))
252262import Data.Functor.Identity (Identity(..))
253263import Data.Proxy (Proxy(..))
@@ -274,13 +284,17 @@ import GHC.Ptr (Ptr(..))
274284import Streamly.Internal.Data.Fold.Type (Fold(..))
275285import Streamly.Internal.Data.Producer.Type (Producer (..))
276286import Streamly.Internal.Data.Stream.Type (Stream)
287+import Streamly.Internal.Data.Parser.Type (Parser (..))
277288import Streamly.Internal.Data.StreamK.Type (StreamK)
278289import Streamly.Internal.Data.SVar.Type (adaptState, defState)
290+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
279291import Streamly.Internal.Data.Unfold.Type (Unfold(..))
280292import Streamly.Internal.System.IO (arrayPayloadSize, defaultChunkSize)
281293
282294import qualified Streamly.Internal.Data.Fold.Type as FL
283295import qualified Streamly.Internal.Data.MutByteArray.Type as Unboxed
296+import qualified Streamly.Internal.Data.Parser.Type as Parser
297+-- import qualified Streamly.Internal.Data.Fold.Type as Fold
284298import qualified Streamly.Internal.Data.Producer as Producer
285299import qualified Streamly.Internal.Data.Stream.Type as D
286300import qualified Streamly.Internal.Data.Stream.Lift as D
@@ -2582,7 +2596,256 @@ byteEq :: MonadIO m => MutArray a -> MutArray a -> m Bool
25822596byteEq arr1 arr2 = fmap (EQ ==) $ byteCmp arr1 arr2
25832597
25842598-------------------------------------------------------------------------------
2585−-- NFData
2599+-- Compact
2600+-------------------------------------------------------------------------------
2601+
2602+-- | Parser @pCompactChunksLE n@ coalesces adjacent arrays in the input stream
2603+-- only if the combined size would be less than or equal to n.
2604+--
2605+-- /Internal/
2606+{-# INLINE_NORMAL pCompactChunksLE #-}
2607+pCompactChunksLE ::
2608+ forall m a. (MonadIO m, Unbox a)
2609+ => Int -> Parser (MutArray a) m (MutArray a)
2610+pCompactChunksLE n = Parser step initial extract
2611+
2612+ where
2613+
2614+ nBytes = n * SIZE_OF(a)
2615+
2616+ functionName = "Streamly.Internal.Data.MutArray.pCompactChunksLE"
2617+
2618+ initial =
2619+ return
2620+ $ if n <= 0
2621+ then error
2622+ $ functionName
2623+ ++ ": the size of arrays ["
2624+ ++ show n ++ "] must be a natural number"
2625+ else Parser.IPartial Nothing
2626+
2627+ step Nothing arr =
2628+ return
2629+ $ let len = byteLength arr
2630+ in if len >= nBytes
2631+ then Parser.Done 0 arr
2632+ else Parser.Partial 0 (Just arr)
2633+ step (Just buf) arr =
2634+ let len = byteLength buf + byteLength arr
2635+ in if len > nBytes
2636+ then return $ Parser.Done 1 buf
2637+ else do
2638+ -- XXX check this in the Nothing case
2639+ buf1 <-
2640+ if byteCapacity buf < nBytes
2641+ then liftIO $ realloc nBytes buf
2642+ else return buf
2643+ buf2 <- spliceUnsafe buf1 arr
2644+ return $ Parser.Partial 0 (Just buf2)
2645+
2646+ extract Nothing = return $ Parser.Done 0 nil
2647+ extract (Just buf) = return $ Parser.Done 0 buf
2648+
2649+data SpliceState s arr
2650+ = SpliceInitial s
2651+ | SpliceBuffering s arr
2652+ | SpliceYielding arr (SpliceState s arr)
2653+ | SpliceFinish
2654+
2655+-- This mutates the first array (if it has space) to append values from the
2656+-- second one. This would work for immutable arrays as well because an
2657+-- immutable array never has additional space so a new array is allocated
2658+-- instead of mutating it.
2659+
2660+-- | Scan @rCompactChunksLE n@ coalesces adjacent arrays in the input stream
2661+-- only if the combined size would be less than or equal to n.
2662+--
2663+-- /Internal/
2664+{-# INLINE_NORMAL rCompactChunksLE #-}
2665+rCompactChunksLE :: (MonadIO m, Unbox a)
2666+ => Int -> D.Stream m (MutArray a) -> D.Stream m (MutArray a)
2667+rCompactChunksLE n (D.Stream step state) =
2668+ D.Stream step' (SpliceInitial state)
2669+
2670+ where
2671+
2672+ functionName = "Streamly.Internal.Data.MutArray.rCompactChunksLE"
2673+
2674+ {-# INLINE_LATE step' #-}
2675+ step' gst (SpliceInitial st) = do
2676+ when (n <= 0) $
2677+ -- XXX we can pass the module string from the higher level API
2678+ error $ functionName ++ ": the size of arrays [" ++ show n
2679+ ++ "] must be a natural number"
2680+ r <- step gst st
2681+ case r of
2682+ D.Yield arr s -> return $
2683+ let len = byteLength arr
2684+ in if len >= n
2685+ then D.Skip (SpliceYielding arr (SpliceInitial s))
2686+ else D.Skip (SpliceBuffering s arr)
2687+ D.Skip s -> return $ D.Skip (SpliceInitial s)
2688+ D.Stop -> return D.Stop
2689+
2690+ step' gst (SpliceBuffering st buf) = do
2691+ r <- step gst st
2692+ case r of
2693+ D.Yield arr s -> do
2694+ let len = byteLength buf + byteLength arr
2695+ if len > n
2696+ then return $
2697+ D.Skip (SpliceYielding buf (SpliceBuffering s arr))
2698+ else do
2699+ -- XXX check this in the SpliceInitial case
2700+ buf1 <- if byteCapacity buf < n
2701+ then liftIO $ realloc n buf
2702+ else return buf
2703+ buf2 <- spliceUnsafe buf1 arr
2704+ return $ D.Skip (SpliceBuffering s buf2)
2705+ D.Skip s -> return $ D.Skip (SpliceBuffering s buf)
2706+ D.Stop -> return $ D.Skip (SpliceYielding buf SpliceFinish)
2707+
2708+ step' _ SpliceFinish = return D.Stop
2709+
2710+ step' _ (SpliceYielding arr next) = return $ D.Yield arr next
2711+
2712+-- | Fold @fCompactChunksGE n@ coalesces adjacent arrays in the input stream
2713+-- until the size becomes greater than or equal to n.
2714+--
2715+{-# INLINE_NORMAL fCompactChunksGE #-}
2716+fCompactChunksGE ::
2717+ forall m a. (MonadIO m, Unbox a)
2718+ => Int -> FL.Fold m (MutArray a) (MutArray a)
2719+fCompactChunksGE n = Fold step initial extract extract
2720+
2721+ where
2722+
2723+ nBytes = n * SIZE_OF(a)
2724+
2725+ functionName = "Streamly.Internal.Data.MutArray.fCompactChunksGE"
2726+
2727+ initial =
2728+ return
2729+ $ if n < 0
2730+ then error
2731+ $ functionName
2732+ ++ ": the size of arrays ["
2733+ ++ show n ++ "] must be a natural number"
2734+ else FL.Partial Nothing
2735+
2736+ step Nothing arr =
2737+ return
2738+ $ let len = byteLength arr
2739+ in if len >= nBytes
2740+ then FL.Done arr
2741+ else FL.Partial (Just arr)
2742+ step (Just buf) arr = do
2743+ let len = byteLength buf + byteLength arr
2744+ buf1 <-
2745+ if byteCapacity buf < len
2746+ then liftIO $ realloc (max len nBytes) buf
2747+ else return buf
2748+ buf2 <- spliceUnsafe buf1 arr
2749+ if len >= n
2750+ then return $ FL.Done buf2
2751+ else return $ FL.Partial (Just buf2)
2752+
2753+ extract Nothing = return nil
2754+ extract (Just buf) = return buf
2755+
2756+-- | Like 'compactChunksGE' but for transforming folds instead of stream.
2757+--
2758+-- >>> lCompactChunksGE n = Fold.many (MutArray.fCompactChunksGE n)
2759+--
2760+{-# INLINE_NORMAL lCompactChunksGE #-}
2761+lCompactChunksGE :: (MonadIO m, Unbox a)
2762+ => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()
2763+-- lCompactChunksGE n = Fold.many (fCompactChunksGE n)
2764+lCompactChunksGE n (Fold step1 initial1 _ final1) =
2765+ Fold step initial extract final
2766+
2767+ where
2768+
2769+ functionName = "Streamly.Internal.Data.MutArray.lCompactChunksGE"
2770+
2771+ initial = do
2772+ when (n <= 0) $
2773+ -- XXX we can pass the module string from the higher level API
2774+ error $ functionName ++ ": the size of arrays ["
2775+ ++ show n ++ "] must be a natural number"
2776+
2777+ r <- initial1
2778+ return $ first (Tuple' Nothing) r
2779+
2780+ step (Tuple' Nothing r1) arr =
2781+ let len = byteLength arr
2782+ in if len >= n
2783+ then do
2784+ r <- step1 r1 arr
2785+ case r of
2786+ FL.Done _ -> return $ FL.Done ()
2787+ FL.Partial s -> do
2788+ _ <- final1 s
2789+ res <- initial1
2790+ return $ first (Tuple' Nothing) res
2791+ else return $ FL.Partial $ Tuple' (Just arr) r1
2792+
2793+ step (Tuple' (Just buf) r1) arr = do
2794+ let len = byteLength buf + byteLength arr
2795+ buf1 <- if byteCapacity buf < len
2796+ then liftIO $ realloc (max n len) buf
2797+ else return buf
2798+ buf2 <- spliceUnsafe buf1 arr
2799+
2800+ -- XXX this is common in both the equations of step
2801+ if len >= n
2802+ then do
2803+ r <- step1 r1 buf2
2804+ case r of
2805+ FL.Done _ -> return $ FL.Done ()
2806+ FL.Partial s -> do
2807+ _ <- final1 s
2808+ res <- initial1
2809+ return $ first (Tuple' Nothing) res
2810+ else return $ FL.Partial $ Tuple' (Just buf2) r1
2811+
2812+ -- XXX Several folds do extract >=> final, therefore, we need to make final
2813+ -- return "m b" rather than using extract post it if we want extract to be
2814+ -- partial.
2815+ --
2816+ -- extract forces the pending buffer to be sent to the fold which is not
2817+ -- what we want.
2818+ extract _ = error "lCompactChunksGE: not designed for scanning"
2819+
2820+ final (Tuple' Nothing r1) = final1 r1
2821+ final (Tuple' (Just buf) r1) = do
2822+ r <- step1 r1 buf
2823+ case r of
2824+ FL.Partial rr -> final1 rr
2825+ FL.Done _ -> return ()
2826+
2827+-- | @compactChunksGE n stream@ coalesces adjacent arrays in the @stream@ until
2828+-- the size becomes greater than or equal to @n@.
2829+--
2830+-- >>> compactChunksGE n = Stream.foldMany (MutArray.fCompactChunksGE n)
2831+--
2832+{-# INLINE compactChunksGE #-}
2833+compactChunksGE ::
2834+ (MonadIO m, Unbox a)
2835+ => Int -> Stream m (MutArray a) -> Stream m (MutArray a)
2836+compactChunksGE n = D.foldMany (fCompactChunksGE n)
2837+
2838+-- | 'compactChunksEQ n' coalesces adajacent arrays in the input stream to
2839+-- arrays of exact size @n@.
2840+--
2841+-- /Unimplemented/
2842+{-# INLINE compactChunksEQ #-}
2843+compactChunksEQ :: -- (MonadIO m, Unbox a) =>
2844+ Int -> Stream m (MutArray a) -> Stream m (MutArray a)
2845+compactChunksEQ _n = undefined -- D.parseManyD (pCompactChunksEQ n)
2846+
2847+-------------------------------------------------------------------------------
2848+-- In-place mutation algorithms
25862849-------------------------------------------------------------------------------
25872850
25882851-- | Strip elements which match with predicate from both ends.