b7d8b96 · streamly · git.iko.soy

git.iko.soy/streamly

Add pinned versions of compaction functions

commitb7d8b96728398e7636338d89bfe041dedbaadfc2 authorHarendra Kumar <harendra@composewell.com> dateMon 8 Jan 2024, 18:42 (3 years ago) parentaf57e26

6 files changed+163−65

compare with master

unifiedsplit

core/src/Streamly/Internal/Data/Array.hs+20 core/src/Streamly/Internal/Data/Array/Type.hs+21−9 core/src/Streamly/Internal/Data/MutArray.hs+20−3 core/src/Streamly/Internal/Data/MutArray/Type.hs+100−50 core/src/Streamly/Internal/Data/MutByteArray/Type.hs+1−1 core/src/Streamly/Internal/FileSystem/Handle.hs+1−2

core/src/Streamly/Internal/Data/Array.hs+20

@@ -74,6 +74,8 @@ module Streamly.Internal.Data.Array

7474 , interposeSuffix

7575 , intercalateSuffix

7676

77+ , compactLE

78+ , pinnedCompactLE

7779 , compactOnByte

7880 , compactOnByteSuffix

7981

@@ -624,6 +626,24 @@ intercalateSuffix :: (Monad m, Unbox a)

624626 => Array a -> Stream m (Array a) -> Stream m a

625627intercalateSuffix = D.intercalateSuffix reader

626628

629+-- | @compactLE n@ coalesces adjacent arrays in the input stream

630+-- only if the combined size would be less than or equal to n.

631+--

632+-- Generates unpinned arrays irrespective of the pinning status of input

633+-- arrays.

634+{-# INLINE_NORMAL compactLE #-}

635+compactLE :: (MonadIO m, Unbox a)

636+ => Int -> Stream m (Array a) -> Stream m (Array a)

637+compactLE n stream =

638+ D.map unsafeFreeze $ MA.compactLE n $ D.map unsafeThaw stream

639+

640+-- | Pinned version of 'compactLE'.

641+{-# INLINE_NORMAL pinnedCompactLE #-}

642+pinnedCompactLE :: (MonadIO m, Unbox a)

643+ => Int -> Stream m (Array a) -> Stream m (Array a)

644+pinnedCompactLE n stream =

645+ D.map unsafeFreeze $ MA.pinnedCompactLE n $ D.map unsafeThaw stream

646+

627647-- | Split a stream of arrays on a given separator byte, dropping the separator

628648-- and coalescing all the arrays between two separators into a single array.

629649--

core/src/Streamly/Internal/Data/Array/Type.hs+21−9

@@ -123,9 +123,10 @@ module Streamly.Internal.Data.Array.Type

123123

124124 -- *** Compact

125125 -- | Append the arrays in a stream to form a stream of larger arrays.

126− , compactLE

127126 , fCompactGE

127+ , fPinnedCompactGE

128128 , lCompactGE

129+ , lPinnedCompactGE

129130 , compactGE

130131

131132 -- ** Deprecated

@@ -502,26 +503,28 @@ flattenArraysRev = concatRev

502503-- arrays would have no capacity to append, therefore, a copy will be forced

503504-- anyway.

504505

505−-- | Scan @compactLE n@ coalesces adjacent arrays in the input stream

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

507−{-# INLINE_NORMAL compactLE #-}

508−compactLE :: (MonadIO m, Unbox a)

509− => Int -> Stream m (Array a) -> Stream m (Array a)

510−compactLE n stream =

511− D.map unsafeFreeze $ MA.rCompactLE n $ D.map unsafeThaw stream

512−

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

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

515508--

509+-- Generates unpinned arrays irrespective of the pinning status of input

510+-- arrays.

516511{-# INLINE_NORMAL fCompactGE #-}

517512fCompactGE :: (MonadIO m, Unbox a) => Int -> Fold m (Array a) (Array a)

518513fCompactGE n = fmap unsafeFreeze $ Fold.lmap unsafeThaw $ MA.fCompactGE n

519514

515+-- | PInned version of 'fCompactGE'.

516+{-# INLINE_NORMAL fPinnedCompactGE #-}

517+fPinnedCompactGE :: (MonadIO m, Unbox a) => Int -> Fold m (Array a) (Array a)

518+fPinnedCompactGE n =

519+ fmap unsafeFreeze $ Fold.lmap unsafeThaw $ MA.fPinnedCompactGE n

520+

520521-- | @compactGE n stream@ coalesces adjacent arrays in the @stream@ until

521522-- the size becomes greater than or equal to @n@.

522523--

523524-- >>> compactGE n = Stream.foldMany (Array.fCompactGE n)

524525--

526+-- Generates unpinned arrays irrespective of the pinning status of input

527+-- arrays.

525528{-# INLINE compactGE #-}

526529compactGE ::

527530 (MonadIO m, Unbox a)

@@ -533,12 +536,21 @@ compactGE n stream =

533536--

534537-- >>> lCompactGE n = Fold.many (Array.fCompactGE n)

535538--

539+-- Generates unpinned arrays irrespective of the pinning status of input

540+-- arrays.

536541{-# INLINE_NORMAL lCompactGE #-}

537542lCompactGE :: (MonadIO m, Unbox a)

538543 => Int -> Fold m (Array a) () -> Fold m (Array a) ()

539544lCompactGE n fld =

540545 Fold.lmap unsafeThaw $ MA.lCompactGE n (Fold.lmap unsafeFreeze fld)

541546

547+-- | Pinned version of 'lCompactGE'.

548+{-# INLINE_NORMAL lPinnedCompactGE #-}

549+lPinnedCompactGE :: (MonadIO m, Unbox a)

550+ => Int -> Fold m (Array a) () -> Fold m (Array a) ()

551+lPinnedCompactGE n fld =

552+ Fold.lmap unsafeThaw $ MA.lPinnedCompactGE n (Fold.lmap unsafeFreeze fld)

553+

542554-------------------------------------------------------------------------------

543555-- Splitting

544556-------------------------------------------------------------------------------

core/src/Streamly/Internal/Data/MutArray.hs+20−3

@@ -14,6 +14,7 @@ module Streamly.Internal.Data.MutArray

1414 , sliceIndexerFromLen

1515 , slicerFromLen

1616 , compactLE

17+ , pinnedCompactLE

1718 , compactOnByte

1819 , compactOnByteSuffix

1920 -- * Unboxed IORef

@@ -29,6 +30,7 @@ where

2930

3031import Control.Monad.IO.Class (MonadIO(..))

3132import Data.Word (Word8)

33+import Streamly.Internal.Data.MutByteArray.Type (PinnedState(..))

3234import Streamly.Internal.Data.Stream.Type (Stream)

3335import Streamly.Internal.Data.Unbox (Unbox)

3436import Streamly.Internal.Data.Unfold.Type (Unfold(..))

@@ -91,15 +93,30 @@ getSlicesFromLen = slicerFromLen

9193-- Compacting Streams of Arrays

9294-------------------------------------------------------------------------------

9395

94−-- | Scan @compactLE n@ coalesces adjacent arrays in the input stream

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

96+-- | @compactLE maxElems@ coalesces adjacent arrays in the input stream

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

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

99+-- larger than maxElems.

100+--

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

102+--

103+-- Generates unpinned arrays irrespective of the pinning status of input

104+-- arrays.

96105{-# INLINE compactLE #-}

97106compactLE :: (MonadIO m, Unbox a) =>

98107 Int -> Stream m (MutArray a) -> Stream m (MutArray a)

99108-- XXX compactLE can be moved to MutArray/Type if we are not using the parser

100109-- to implement it.

110+compactLE = compactLeAs Unpinned

111+-- The parser version turns out to be a little bit slower.

101112-- compactLE n = Stream.catRights . Stream.parseManyD (pCompactLE n)

102−compactLE = rCompactLE

113+

114+-- | Pinned version of 'compactLE'.

115+{-# INLINE pinnedCompactLE #-}

116+pinnedCompactLE :: forall m a. (MonadIO m, Unbox a)

117+ => Int -> Stream m (MutArray a) -> Stream m (MutArray a)

118+pinnedCompactLE = compactLeAs Pinned

119+-- pinnedCompactLE n = Stream.catRights . Stream.parseManyD (pPinnedCompactLE n)

103120

104121data SplitState s arr

105122 = Initial s

core/src/Streamly/Internal/Data/MutArray/Type.hs+100−50

@@ -251,9 +251,12 @@ module Streamly.Internal.Data.MutArray.Type

251251 -- | Append the arrays in a stream to form a stream of larger arrays.

252252 , SpliceState (..)

253253 , pCompactLE

254− , rCompactLE

254+ , pPinnedCompactLE

255+ , compactLeAs

255256 , fCompactGE

257+ , fPinnedCompactGE

256258 , lCompactGE

259+ , lPinnedCompactGE

257260 , compactGE

258261 , compactEQ

259262

@@ -839,15 +842,15 @@ roundDownTo elemSize size = size - (size `mod` elemSize)

839842-- Since this is not inlined Unboxed consrraint leads to dictionary passing

840843-- which complicates some inspection tests.

841844--

842−{-# NOINLINE reallocExplicit #-}

843−reallocExplicit :: Int -> Int -> MutArray a -> IO (MutArray a)

844−reallocExplicit elemSize newCapacityInBytes MutArray{..} = do

845+{-# NOINLINE reallocExplicitAs #-}

846+reallocExplicitAs :: PinnedState -> Int -> Int -> MutArray a -> IO (MutArray a)

847+reallocExplicitAs ps elemSize newCapacityInBytes MutArray{..} = do

845848 assertM(arrEnd <= arrBound)

846849

847850 -- Allocate new array

848851 let newCapMaxInBytes = roundUpLargeArray newCapacityInBytes

849852 contents <-

850− if Unboxed.isPinned arrContents

853+ if ps == Pinned

851854 then Unboxed.pinnedNew newCapMaxInBytes

852855 else Unboxed.new newCapMaxInBytes

853856 let !(MutByteArray mbarrFrom#) = arrContents

@@ -887,7 +890,12 @@ reallocExplicit elemSize newCapacityInBytes MutArray{..} = do

887890-- If the original array is pinned, the newly allocated array is also pinned.

888891{-# INLINABLE realloc #-}

889892realloc :: forall m a. (MonadIO m, Unbox a) => Int -> MutArray a -> m (MutArray a)

890−realloc bytes arr = liftIO $ reallocExplicit (SIZE_OF(a)) bytes arr

893+realloc bytes arr =

894+ let ps =

895+ if isPinned arr

896+ then Pinned

897+ else Unpinned

898+ in liftIO $ reallocExplicitAs ps (SIZE_OF(a)) bytes arr

891899

892900-- | @reallocWith label capSizer minIncrBytes array@. The label is used

893901-- in error messages and the capSizer is used to determine the capacity of the

@@ -2284,7 +2292,7 @@ writeWithAs ps elemCount =

22842292 | INDEX_NEXT(end,a) > bound = do

22852293 let oldSize = end - start

22862294 newSize = max (oldSize * 2) 1

2287− arr1 <- liftIO $ reallocExplicit (SIZE_OF(a)) newSize arr

2295+ arr1 <- liftIO $ reallocExplicitAs ps (SIZE_OF(a)) newSize arr

22882296 snocUnsafe arr1 x

22892297 step arr x = snocUnsafe arr x

22902298

@@ -2943,19 +2951,11 @@ byteEq arr1 arr2 = fmap (EQ ==) $ byteCmp arr1 arr2

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

29442952-- exactly in the space.

29452953

2946−-- | Parser @pCompactLE maxElems@ coalesces adjacent arrays in the input stream

2947−-- only if the combined size would be less than or equal to @maxElems@

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

2949−-- larger than maxElems.

2950−--

2951−-- @maxElems@ must be greater than 0.

2952−--

2953−-- /Internal/

2954−{-# INLINE_NORMAL pCompactLE #-}

2955−pCompactLE ::

2954+{-# INLINE_NORMAL pCompactLeAs #-}

2955+pCompactLeAs ::

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

2957− => Int -> Parser (MutArray a) m (MutArray a)

2958−pCompactLE maxElems = Parser step initial extract

2957+ => PinnedState -> Int -> Parser (MutArray a) m (MutArray a)

2958+pCompactLeAs ps maxElems = Parser step initial extract

29592959

29602960 where

29612961

@@ -2986,7 +2986,8 @@ pCompactLE maxElems = Parser step initial extract

29862986 else do

29872987 buf1 <-

29882988 if byteCapacity buf < maxBytes

2989− then realloc maxBytes buf

2989+ then liftIO $ reallocExplicitAs

2990+ ps (SIZE_OF(a)) maxBytes buf

29902991 else return buf

29912992 buf2 <- spliceUnsafe buf1 arr

29922993 return $ Parser.Partial 0 (Just buf2)

@@ -2994,6 +2995,30 @@ pCompactLE maxElems = Parser step initial extract

29942995 extract Nothing = return $ Parser.Done 0 nil

29952996 extract (Just buf) = return $ Parser.Done 0 buf

29962997

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

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

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

3001+-- larger than maxElems.

3002+--

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

3004+--

3005+-- Generates unpinned arrays irrespective of the pinning status of input

3006+-- arrays.

3007+--

3008+-- /Internal/

3009+{-# INLINE pCompactLE #-}

3010+pCompactLE ::

3011+ forall m a. (MonadIO m, Unbox a)

3012+ => Int -> Parser (MutArray a) m (MutArray a)

3013+pCompactLE = pCompactLeAs Unpinned

3014+

3015+-- | Pinned version of 'pCompactLE'.

3016+{-# INLINE pPinnedCompactLE #-}

3017+pPinnedCompactLE ::

3018+ forall m a. (MonadIO m, Unbox a)

3019+ => Int -> Parser (MutArray a) m (MutArray a)

3020+pPinnedCompactLE = pCompactLeAs Pinned

3021+

29973022data SpliceState s arr

29983023 = SpliceInitial s

29993024 | SpliceBuffering s arr

@@ -3005,18 +3030,10 @@ data SpliceState s arr

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

30063031-- instead of mutating it.

30073032

3008−-- | Scan @rCompactLE maxElems@ coalesces adjacent arrays in the input stream

3009−-- only if the combined size would be less than or equal to @maxElems@

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

3011−-- larger than maxElems.

3012−--

3013−-- @maxElems@ must be greater than 0.

3014−--

3015−-- /Internal/

3016−{-# INLINE_NORMAL rCompactLE #-}

3017−rCompactLE :: forall m a. (MonadIO m, Unbox a)

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

3019−rCompactLE maxElems (D.Stream step state) =

3033+{-# INLINE_NORMAL compactLeAs #-}

3034+compactLeAs :: forall m a. (MonadIO m, Unbox a)

3035+ => PinnedState -> Int -> D.Stream m (MutArray a) -> D.Stream m (MutArray a)

3036+compactLeAs ps maxElems (D.Stream step state) =

30203037 D.Stream step' (SpliceInitial state)

30213038

30223039 where

@@ -3052,7 +3069,8 @@ rCompactLE maxElems (D.Stream step state) =

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

30533070 else do

30543071 buf1 <- if byteCapacity buf < maxBytes

3055− then realloc maxBytes buf

3072+ then liftIO $ reallocExplicitAs

3073+ ps (SIZE_OF(a)) maxBytes buf

30563074 else return buf

30573075 buf2 <- spliceUnsafe buf1 arr

30583076 return $ D.Skip (SpliceBuffering s buf2)

@@ -3063,14 +3081,12 @@ rCompactLE maxElems (D.Stream step state) =

30633081

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

30653083

3066−-- | Fold @fCompactGE minElems@ coalesces adjacent arrays in the input stream

3067−-- until the size becomes greater than or equal to @minElems@.

3068−--

3069−{-# INLINE_NORMAL fCompactGE #-}

3070−fCompactGE ::

3084+

3085+{-# INLINE_NORMAL fCompactGeAs #-}

3086+fCompactGeAs ::

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

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

3073−fCompactGE minElems = Fold step initial extract extract

3088+ => PinnedState -> Int -> FL.Fold m (MutArray a) (MutArray a)

3089+fCompactGeAs ps minElems = Fold step initial extract extract

30743090

30753091 where

30763092

@@ -3098,7 +3114,8 @@ fCompactGE minElems = Fold step initial extract extract

30983114 let len = byteLength buf + byteLength arr

30993115 buf1 <-

31003116 if byteCapacity buf < len

3101− then realloc (max len minBytes) buf

3117+ then liftIO $ reallocExplicitAs

3118+ ps (SIZE_OF(a)) (max minBytes len) buf

31023119 else return buf

31033120 buf2 <- spliceUnsafe buf1 arr

31043121 if len >= minBytes

@@ -3108,15 +3125,30 @@ fCompactGE minElems = Fold step initial extract extract

31083125 extract Nothing = return nil

31093126 extract (Just buf) = return buf

31103127

3111−-- | Like 'compactGE' but for transforming folds instead of stream.

3112−--

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

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

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

31143130--

3115−{-# INLINE_NORMAL lCompactGE #-}

3116−lCompactGE :: forall m a. (MonadIO m, Unbox a)

3117− => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()

3118−-- lCompactGE n = Fold.many (fCompactGE n)

3119−lCompactGE minElems (Fold step1 initial1 _ final1) =

3131+-- Generates unpinned arrays irrespective of the pinning status of input

3132+-- arrays.

3133+{-# INLINE fCompactGE #-}

3134+fCompactGE ::

3135+ forall m a. (MonadIO m, Unbox a)

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

3137+fCompactGE = fCompactGeAs Unpinned

3138+

3139+-- | Pinned version of 'fCompactGE'.

3140+{-# INLINE fPinnedCompactGE #-}

3141+fPinnedCompactGE ::

3142+ forall m a. (MonadIO m, Unbox a)

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

3144+fPinnedCompactGE = fCompactGeAs Pinned

3145+

3146+{-# INLINE_NORMAL lCompactGeAs #-}

3147+lCompactGeAs :: forall m a. (MonadIO m, Unbox a)

3148+ => PinnedState -> Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()

3149+-- The fold version turns out to be a little bit slower.

3150+-- lCompactGeAs ps n = FL.many (fCompactGeAs ps n)

3151+lCompactGeAs ps minElems (Fold step1 initial1 _ final1) =

31203152 Fold step initial extract final

31213153

31223154 where

@@ -3154,7 +3186,8 @@ lCompactGE minElems (Fold step1 initial1 _ final1) =

31543186 step (Tuple' (Just buf) r1) arr = do

31553187 let len = byteLength buf + byteLength arr

31563188 buf1 <- if byteCapacity buf < len

3157− then realloc (max minBytes len) buf

3189+ then liftIO $ reallocExplicitAs

3190+ ps (SIZE_OF(a)) (max minBytes len) buf

31583191 else return buf

31593192 buf2 <- spliceUnsafe buf1 arr

31603193 runInner len r1 buf2

@@ -3174,6 +3207,23 @@ lCompactGE minElems (Fold step1 initial1 _ final1) =

31743207 FL.Partial rr -> final1 rr

31753208 FL.Done _ -> return ()

31763209

3210+-- | Like 'compactGE' but for transforming folds instead of stream.

3211+--

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

3213+--

3214+-- Generates unpinned arrays irrespective of the pinning status of input

3215+-- arrays.

3216+{-# INLINE lCompactGE #-}

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

3218+ => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()

3219+lCompactGE = lCompactGeAs Unpinned

3220+

3221+-- | Pinned version of 'lCompactGE'.

3222+{-# INLINE lPinnedCompactGE #-}

3223+lPinnedCompactGE :: forall m a. (MonadIO m, Unbox a)

3224+ => Int -> Fold m (MutArray a) () -> Fold m (MutArray a) ()

3225+lPinnedCompactGE = lCompactGeAs Pinned

3226+

31773227-- | @compactGE n stream@ coalesces adjacent arrays in the @stream@ until

31783228-- the size becomes greater than or equal to @n@.

31793229--

core/src/Streamly/Internal/Data/MutByteArray/Type.hs+1−1

@@ -58,7 +58,7 @@ import GHC.Exts

5858

5959data PinnedState

6060 = Pinned

61− | Unpinned

61+ | Unpinned deriving (Show, Eq)

6262

6363-- XXX can use UnliftedNewtypes

6464

core/src/Streamly/Internal/FileSystem/Handle.hs+1−2

@@ -137,8 +137,7 @@ import Streamly.Internal.Data.Stream.Type (Stream)

137137import Streamly.Internal.System.IO (defaultChunkSize)

138138

139139import qualified Streamly.Data.Fold as FL

140−import qualified Streamly.Data.Array as A

141−import qualified Streamly.Internal.Data.Array.Type as A

140+import qualified Streamly.Internal.Data.Array as A

142141import qualified Streamly.Internal.Data.MutArray.Type as MArray

143142import qualified Streamly.Internal.Data.Refold.Type as Refold

144143import qualified Streamly.Internal.Data.Fold.Type as FL(refoldMany)