Add createOfLast in MutArray module and update Array.createOfLast

composewell · Dec 11, 2024 · 36cfebc · 36cfebc
1 parent 2761290
commit 36cfebc
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/core/src/Streamly/Internal/Data/Array.hs b/core/src/Streamly/Internal/Data/Array.hs
@@ -135,7 +135,6 @@ import qualified Streamly.Internal.Data.Fold.Type as Fold
 import qualified Streamly.Internal.Data.Serialize.Type as Serialize
 import qualified Streamly.Internal.Data.MutByteArray.Type as MBA
 import qualified Streamly.Internal.Data.MutArray as MA
-import qualified Streamly.Internal.Data.RingArray as RB
 import qualified Streamly.Internal.Data.Parser as Parser
 -- import qualified Streamly.Internal.Data.ParserK as ParserK
 import qualified Streamly.Internal.Data.Stream as D
@@ -214,15 +213,8 @@ last = getIndexRev 0
 -- stream to an 'Array'.
 --
 {-# INLINE createOfLast #-}
-createOfLast ::
-       (Unbox a, MonadIO m) => Int -> Fold m a (Array a)
-createOfLast n = Fold.ifThen (pure (n <= 0)) (Fold.fromPure empty) lst
-
-    where
-
-    lst =
-        let f = fmap unsafeFreeze . RB.toMutArray
-         in Fold.rmapM f $ RB.createOfLast n
+createOfLast :: (Unbox a, MonadIO m) => Int -> Fold m a (Array a)
+createOfLast = fmap unsafeFreeze . MA.createOfLast
 
 {-# DEPRECATED writeLastN "Please use createOfLast instead." #-}
 {-# INLINE writeLastN #-}

diff --git 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)