e135f97 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Refactor array compaction functions

e135f97Harendra Kumar3 years ago

core/docs/Changelog.md+2

compare with master

unifiedsplit

@@ -5,6 +5,8 @@

5* Add fixities `infixr 5` for `cons` and `consM` functions.

6* Fix a bug in Array `Eq` instance when the type is a sum type with

7 differently sized constructors.

8+* lpackArraysChunksOf, compact, writeChunksWith, putChunksWith now take the

9+ buffer size in number of array elements instead of bytes.

9## 0.2.1 (Dec 2023)

11## 0.2.1 (Dec 2023)

core/src/Streamly/Internal/Data/MutArray/Type.hs+71−60

@@ -2635,46 +2635,55 @@ byteEq arr1 arr2 = fmap (EQ ==) $ byteCmp arr1 arr2

2635-- Compact

2636-------------------------------------------------------------------------------

2637

2638+-- Note: LE versions avoid an extra copy compared to GE. LE parser trades

2639+-- backtracking one array in lieu of avoiding a copy. However, LE and GE both

2640+-- can leave some memory unused. They can split the last array to fit it

2641+-- exactly in the space.

2642+

2638−-- | Parser @pCompactLE n@ coalesces adjacent arrays in the input stream

2643+-- | Parser @pCompactLE maxElems@ coalesces adjacent arrays in the input stream

2639−-- only if the combined size would be less than or equal to n.

2644+-- only if the combined size would be less than or equal to @maxElems@

2645+-- elements. Note that it won't split an array if the original array is already

2646+-- larger than maxElems.

2647+--

2648+-- @maxElems@ must be greater than 0.

2640--

2649--

2641-- /Internal/

2650-- /Internal/

2642{-# INLINE_NORMAL pCompactLE #-}

2651{-# INLINE_NORMAL pCompactLE #-}

2643pCompactLE ::

2652pCompactLE ::

2644 forall m a. (MonadIO m, Unbox a)

2653 forall m a. (MonadIO m, Unbox a)

2645 => Int -> Parser (MutArray a) m (MutArray a)

2654 => Int -> Parser (MutArray a) m (MutArray a)

2646−pCompactLE n = Parser step initial extract

2655+pCompactLE maxElems = Parser step initial extract

2647

2656

2648 where

2657 where

2649

2658

2650− nBytes = n * SIZE_OF(a)

2659+ maxBytes = maxElems * SIZE_OF(a)

2651

2660

2652 functionName = "Streamly.Internal.Data.MutArray.pCompactLE"

2661 functionName = "Streamly.Internal.Data.MutArray.pCompactLE"

2653

2662

2654 initial =

2663 initial =

2655 return

2664 return

2656− $ if n <= 0

2665+ $ if maxElems <= 0

2657 then error

2666 then error

2658 $ functionName

2667 $ functionName

2659 ++ ": the size of arrays ["

2668 ++ ": the size of arrays ["

2660− ++ show n ++ "] must be a natural number"

2669+ ++ show maxElems ++ "] must be a natural number"

2661 else Parser.IPartial Nothing

2670 else Parser.IPartial Nothing

2662

2671

2663 step Nothing arr =

2672 step Nothing arr =

2664 return

2673 return

2665 $ let len = byteLength arr

2674 $ let len = byteLength arr

2666− in if len >= nBytes

2675+ in if len >= maxBytes

2667 then Parser.Done 0 arr

2676 then Parser.Done 0 arr

2668 else Parser.Partial 0 (Just arr)

2677 else Parser.Partial 0 (Just arr)

2678+ -- XXX Split the last array to use the space more compactly.

2669 step (Just buf) arr =

2679 step (Just buf) arr =

2670 let len = byteLength buf + byteLength arr

2680 let len = byteLength buf + byteLength arr

2671− in if len > nBytes

2681+ in if len > maxBytes

2672 then return $ Parser.Done 1 buf

2682 then return $ Parser.Done 1 buf

2673 else do

2683 else do

2674− -- XXX check this in the Nothing case

2675 buf1 <-

2684 buf1 <-

2676− if byteCapacity buf < nBytes

2685+ if byteCapacity buf < maxBytes

2677− then liftIO $ realloc nBytes buf

2686+ then realloc maxBytes buf

2678 else return buf

2687 else return buf

2679 buf2 <- spliceUnsafe buf1 arr

2688 buf2 <- spliceUnsafe buf1 arr

2680 return $ Parser.Partial 0 (Just buf2)

2689 return $ Parser.Partial 0 (Just buf2)

@@ -2693,48 +2702,54 @@ data SpliceState s arr

2693-- immutable array never has additional space so a new array is allocated

2702-- immutable array never has additional space so a new array is allocated

2694-- instead of mutating it.

2703-- instead of mutating it.

2695

2704

2696−-- | Scan @rCompactLE n@ coalesces adjacent arrays in the input stream

2705+-- | Scan @rCompactLE maxElems@ coalesces adjacent arrays in the input stream

2697−-- only if the combined size would be less than or equal to n.

2706+-- only if the combined size would be less than or equal to @maxElems@

2707+-- elements. Note that it won't split an array if the original array is already

2708+-- larger than maxElems.

2709+--

2710+-- @maxElems@ must be greater than 0.

2698--

2711--

2699-- /Internal/

2712-- /Internal/

2700{-# INLINE_NORMAL rCompactLE #-}

2713{-# INLINE_NORMAL rCompactLE #-}

2701−rCompactLE :: (MonadIO m, Unbox a)

2714+rCompactLE :: forall m a. (MonadIO m, Unbox a)

2702 => Int -> D.Stream m (MutArray a) -> D.Stream m (MutArray a)

2715 => Int -> D.Stream m (MutArray a) -> D.Stream m (MutArray a)

2703−rCompactLE n (D.Stream step state) =

2716+rCompactLE maxElems (D.Stream step state) =

2704 D.Stream step' (SpliceInitial state)

2717 D.Stream step' (SpliceInitial state)

2705

2718

2706 where

2719 where

2707

2720

2721+ maxBytes = maxElems * SIZE_OF(a)

2722+

2708 functionName = "Streamly.Internal.Data.MutArray.rCompactLE"

2723 functionName = "Streamly.Internal.Data.MutArray.rCompactLE"

2709

2724

2710 {-# INLINE_LATE step' #-}

2725 {-# INLINE_LATE step' #-}

2711 step' gst (SpliceInitial st) = do

2726 step' gst (SpliceInitial st) = do

2712− when (n <= 0) $

2727+ when (maxElems <= 0) $

2713 -- XXX we can pass the module string from the higher level API

2728 -- XXX we can pass the module string from the higher level API

2714− error $ functionName ++ ": the size of arrays [" ++ show n

2729+ error $ functionName ++ ": the size of arrays [" ++ show maxElems

2715 ++ "] must be a natural number"

2730 ++ "] must be a natural number"

2716 r <- step gst st

2731 r <- step gst st

2717 case r of

2732 case r of

2718 D.Yield arr s -> return $

2733 D.Yield arr s -> return $

2719 let len = byteLength arr

2734 let len = byteLength arr

2720− in if len >= n

2735+ in if len >= maxBytes

2721 then D.Skip (SpliceYielding arr (SpliceInitial s))

2736 then D.Skip (SpliceYielding arr (SpliceInitial s))

2722 else D.Skip (SpliceBuffering s arr)

2737 else D.Skip (SpliceBuffering s arr)

2723 D.Skip s -> return $ D.Skip (SpliceInitial s)

2738 D.Skip s -> return $ D.Skip (SpliceInitial s)

2724 D.Stop -> return D.Stop

2739 D.Stop -> return D.Stop

2725

2740

2741+ -- XXX Split the last array to use the space more compactly.

2726 step' gst (SpliceBuffering st buf) = do

2742 step' gst (SpliceBuffering st buf) = do

2727 r <- step gst st

2743 r <- step gst st

2728 case r of

2744 case r of

2729 D.Yield arr s -> do

2745 D.Yield arr s -> do

2730 let len = byteLength buf + byteLength arr

2746 let len = byteLength buf + byteLength arr

2731− if len > n

2747+ if len > maxBytes

2732 then return $

2748 then return $

2733 D.Skip (SpliceYielding buf (SpliceBuffering s arr))

2749 D.Skip (SpliceYielding buf (SpliceBuffering s arr))

2734 else do

2750 else do

2735− -- XXX check this in the SpliceInitial case

2736− buf1 <- if byteCapacity buf < n

2751+ buf1 <- if byteCapacity buf < maxBytes

2737− then liftIO $ realloc n buf

2752+ then realloc maxBytes buf

2738 else return buf

2753 else return buf

2739 buf2 <- spliceUnsafe buf1 arr

2754 buf2 <- spliceUnsafe buf1 arr

2740 return $ D.Skip (SpliceBuffering s buf2)

2755 return $ D.Skip (SpliceBuffering s buf2)

@@ -2745,44 +2760,45 @@ rCompactLE n (D.Stream step state) =

2745

2760

2746 step' _ (SpliceYielding arr next) = return $ D.Yield arr next

2761 step' _ (SpliceYielding arr next) = return $ D.Yield arr next

2747

2762

2748−-- | Fold @fCompactGE n@ coalesces adjacent arrays in the input stream

2763+-- | Fold @fCompactGE minElems@ coalesces adjacent arrays in the input stream

2749−-- until the size becomes greater than or equal to n.

2764+-- until the size becomes greater than or equal to @minElems@.

2750--

2765--

2751{-# INLINE_NORMAL fCompactGE #-}

2766{-# INLINE_NORMAL fCompactGE #-}

2752fCompactGE ::

2767fCompactGE ::

2753 forall m a. (MonadIO m, Unbox a)

2768 forall m a. (MonadIO m, Unbox a)

2754 => Int -> FL.Fold m (MutArray a) (MutArray a)

2769 => Int -> FL.Fold m (MutArray a) (MutArray a)

2755−fCompactGE n = Fold step initial extract extract

2770+fCompactGE minElems = Fold step initial extract extract

2756

2771

2757 where

2772 where

2758

2773

2759− nBytes = n * SIZE_OF(a)

2774+ minBytes = minElems * SIZE_OF(a)

2760

2775

2761 functionName = "Streamly.Internal.Data.MutArray.fCompactGE"

2776 functionName = "Streamly.Internal.Data.MutArray.fCompactGE"

2762

2777

2763 initial =

2778 initial =

2764 return

2779 return

2765− $ if n < 0

2780+ $ if minElems < 0

2766 then error

2781 then error

2767 $ functionName

2782 $ functionName

2768 ++ ": the size of arrays ["

2783 ++ ": the size of arrays ["

2769− ++ show n ++ "] must be a natural number"

2784+ ++ show minElems ++ "] must be a natural number"

2770 else FL.Partial Nothing

2785 else FL.Partial Nothing

2771

2786

2772 step Nothing arr =

2787 step Nothing arr =

2773 return

2788 return

2774 $ let len = byteLength arr

2789 $ let len = byteLength arr

2775− in if len >= nBytes

2790+ in if len >= minBytes

2776 then FL.Done arr

2791 then FL.Done arr

2777 else FL.Partial (Just arr)

2792 else FL.Partial (Just arr)

2793+ -- XXX Buffer arrays as a list to avoid copy and reallocations

2778 step (Just buf) arr = do

2794 step (Just buf) arr = do

2779 let len = byteLength buf + byteLength arr

2795 let len = byteLength buf + byteLength arr

2780 buf1 <-

2796 buf1 <-

2781 if byteCapacity buf < len

2797 if byteCapacity buf < len

2782− then liftIO $ realloc (max len nBytes) buf

2798+ then realloc (max len minBytes) buf

2783 else return buf

2799 else return buf

2784 buf2 <- spliceUnsafe buf1 arr

2800 buf2 <- spliceUnsafe buf1 arr

2785− if len >= n

2801+ if len >= minBytes

2786 then return $ FL.Done buf2

2802 then return $ FL.Done buf2

2787 else return $ FL.Partial (Just buf2)

2803 else return $ FL.Partial (Just buf2)

2788

2804

@@ -2794,59 +2810,54 @@ fCompactGE n = Fold step initial extract extract

2794-- >>> lCompactGE n = Fold.many (MutArray.fCompactGE n)

2810-- >>> lCompactGE n = Fold.many (MutArray.fCompactGE n)

2795--

2811--

2796{-# INLINE_NORMAL lCompactGE #-}

2812{-# INLINE_NORMAL lCompactGE #-}

2797−lCompactGE :: (MonadIO m, Unbox a)

2813+lCompactGE :: forall m a. (MonadIO m, Unbox a)

2798 => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()

2814 => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()

2799-- lCompactGE n = Fold.many (fCompactGE n)

2815-- lCompactGE n = Fold.many (fCompactGE n)

2800−lCompactGE n (Fold step1 initial1 _ final1) =

2816+lCompactGE minElems (Fold step1 initial1 _ final1) =

2801 Fold step initial extract final

2817 Fold step initial extract final

2802

2818

2803 where

2819 where

2804

2820

2821+ minBytes = minElems * SIZE_OF(a)

2822+

2805 functionName = "Streamly.Internal.Data.MutArray.lCompactGE"

2823 functionName = "Streamly.Internal.Data.MutArray.lCompactGE"

2806

2824

2807 initial = do

2825 initial = do

2808− when (n <= 0) $

2826+ when (minElems <= 0) $

2809 -- XXX we can pass the module string from the higher level API

2827 -- XXX we can pass the module string from the higher level API

2810 error $ functionName ++ ": the size of arrays ["

2828 error $ functionName ++ ": the size of arrays ["

2811− ++ show n ++ "] must be a natural number"

2829+ ++ show minElems ++ "] must be a natural number"

2812

2830

2813 r <- initial1

2831 r <- initial1

2814 return $ first (Tuple' Nothing) r

2832 return $ first (Tuple' Nothing) r

2815

2833

2816− step (Tuple' Nothing r1) arr =

2834+ {-# INLINE runInner #-}

2817− let len = byteLength arr

2835+ runInner len acc buf =

2818− in if len >= n

2819− then do

2820− r <- step1 r1 arr

2821− case r of

2822− FL.Done _ -> return $ FL.Done ()

2823− FL.Partial s -> do

2824− _ <- final1 s

2825− res <- initial1

2826− return $ first (Tuple' Nothing) res

2827− else return $ FL.Partial $ Tuple' (Just arr) r1

2828−

2829− step (Tuple' (Just buf) r1) arr = do

2830− let len = byteLength buf + byteLength arr

2831− buf1 <- if byteCapacity buf < len

2832− then liftIO $ realloc (max n len) buf

2833− else return buf

2834− buf2 <- spliceUnsafe buf1 arr

2835−

2836− -- XXX this is common in both the equations of step

2837− if len >= n

2836+ if len >= minBytes

2838 then do

2837 then do

2839− r <- step1 r1 buf2

2838+ r <- step1 acc buf

2840 case r of

2839 case r of

2841 FL.Done _ -> return $ FL.Done ()

2840 FL.Done _ -> return $ FL.Done ()

2842 FL.Partial s -> do

2841 FL.Partial s -> do

2843 _ <- final1 s

2842 _ <- final1 s

2844 res <- initial1

2843 res <- initial1

2845 return $ first (Tuple' Nothing) res

2844 return $ first (Tuple' Nothing) res

2846− else return $ FL.Partial $ Tuple' (Just buf2) r1

2845+ else return $ FL.Partial $ Tuple' (Just buf) acc

2846+

2847+ step (Tuple' Nothing r1) arr =

2848+ runInner (byteLength arr) r1 arr

2849+

2850+ -- XXX Buffer arrays as a list to avoid copy and reallocations

2851+ step (Tuple' (Just buf) r1) arr = do

2852+ let len = byteLength buf + byteLength arr

2853+ buf1 <- if byteCapacity buf < len

2854+ then realloc (max minBytes len) buf

2855+ else return buf

2856+ buf2 <- spliceUnsafe buf1 arr

2857+ runInner len r1 buf2

2847

2858

2848 -- XXX Several folds do extract >=> final, therefore, we need to make final

2859 -- XXX Several folds do extract >=> final, therefore, we need to make final

2849− -- return "m b" rather than using extract post it if we want extract to be

2860+ -- return "m b" rather than using extract post it if we want extract to be

2850 -- partial.

2861 -- partial.

2851 --

2862 --

2852 -- extract forces the pending buffer to be sent to the fold which is not

2863 -- extract forces the pending buffer to be sent to the fold which is not