diff --git a/core/src/Streamly/Internal/Data/Array.hs b/core/src/Streamly/Internal/Data/Array.hs index 29143e86ca..779bb44997 100644 --- a/core/src/Streamly/Internal/Data/Array.hs +++ b/core/src/Streamly/Internal/Data/Array.hs @@ -210,6 +210,10 @@ last = getIndexRev 0 -- Folds with Array as the container ------------------------------------------------------------------------------- +-- NOTE: We could possible write this in terms of "MutArray.createOfLast" but +-- this causes regression. This is probably because mapping inside "Fold.ifThen" +-- is more efficient than mapping over "Fold.ifTen". +-- -- | @createOfLast n@ folds a maximum of @n@ elements from the end of the input -- stream to an 'Array'. -- diff --git a/core/src/Streamly/Internal/Data/MutArray.hs b/core/src/Streamly/Internal/Data/MutArray.hs index 3c33dd1ec7..fcdbbcb777 100644 --- a/core/src/Streamly/Internal/Data/MutArray.hs +++ b/core/src/Streamly/Internal/Data/MutArray.hs @@ -25,6 +25,7 @@ module Streamly.Internal.Data.MutArray , compactSepByByte_ , compactEndByByte_ , compactEndByLn_ + , createOfLast -- * Unboxed IORef , module Streamly.Internal.Data.IORef.Unboxed @@ -57,10 +58,13 @@ import Streamly.Internal.Data.Serialize.Type (Serialize) import Streamly.Internal.Data.Stream.Type (Stream) import Streamly.Internal.Data.Unbox (Unbox) import Streamly.Internal.Data.Unfold.Type (Unfold(..)) +import Streamly.Internal.Data.Fold.Type (Fold) +import qualified Streamly.Internal.Data.RingArray as RingArray import qualified Streamly.Internal.Data.Serialize.Type as Serialize import qualified Streamly.Internal.Data.Stream.Nesting as Stream import qualified Streamly.Internal.Data.Stream.Type as Stream +import qualified Streamly.Internal.Data.Fold.Type as Fold -- import qualified Streamly.Internal.Data.Stream.Transform as Stream import qualified Streamly.Internal.Data.Unfold as Unfold @@ -337,3 +341,14 @@ compactEndByLn_ :: MonadIO m => Stream m (MutArray Word8) -> Stream m (MutArray Word8) compactEndByLn_ = compactEndByByte_ 10 + +-- | @createOfLast n@ folds a maximum of @n@ elements from the end of the input +-- stream to an 'MutArray'. +-- +{-# INLINE createOfLast #-} +createOfLast :: (Unbox a, MonadIO m) => Int -> Fold m a (MutArray a) +createOfLast n = + Fold.ifThen + (pure (n <= 0)) + (Fold.fromPure empty) + (Fold.rmapM RingArray.toMutArray $ RingArray.createOfLast n)