git.iko.soy/streamly
on master help sign in
Move lastN from StreamD to Array.Storable.Foreign
commit74be5ad9a3d47476d298e9041e7f1b7cbba1cdc2 authorAdithya Kumar <adithya@composewell.com> dateThu 10 Dec 2020, 09:50 (6 years ago) parente44e5b0
2 files changed+30−27
compare with master
unifiedsplit
src/Streamly/Internal/Data/Array/Storable/Foreign.hs+30−3 src/Streamly/Internal/Data/Stream/StreamD.hs−24
src/Streamly/Internal/Data/Array/Storable/Foreign.hs+30−3
@@ -135,7 +135,7 @@ module Streamly.Internal.Data.Array.Storable.Foreign
135135 , fold
136136
137137 -- * Folds with Array as the container
138− , D.lastN
138+ , lastN
139139 )
140140where
141141
@@ -157,17 +157,21 @@ import GHC.Ptr (Ptr(..))
157157import GHC.Prim (touch#)
158158import GHC.IO (IO(..))
159159
160−import Streamly.Internal.Data.Fold.Types (Fold(..))
161−import Streamly.Internal.Data.Unfold.Types (Unfold(..))
162160import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..), length)
161+import Streamly.Internal.Data.Fold.Types (Fold(..))
163162import Streamly.Internal.Data.Stream.Serial (SerialT)
164163import Streamly.Internal.Data.Stream.StreamK.Type (IsStream)
164+import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
165+import Streamly.Internal.Data.Unfold.Types (Unfold(..))
165166
167+import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
166168import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
169+import qualified Streamly.Internal.Data.Fold as FL
167170import qualified Streamly.Internal.Data.Stream.Prelude as P
168171import qualified Streamly.Internal.Data.Stream.Serial as Serial
169172import qualified Streamly.Internal.Data.Stream.StreamD as D
170173import qualified Streamly.Internal.Data.Stream.StreamK as K
174+import qualified Streamly.Memory.Ring as RB
171175
172176-------------------------------------------------------------------------------
173177-- Construction
@@ -301,6 +305,29 @@ null arr = length arr == 0
301305last :: Storable a => Array a -> Maybe a
302306last arr = readIndex arr (length arr - 1)
303307
308+-------------------------------------------------------------------------------
309+-- Folds with Array as the container
310+-------------------------------------------------------------------------------
311+
312+-- | Take last 'n' elements from the stream and discard the rest.
313+{-# INLINE lastN #-}
314+lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
315+lastN n
316+ | n <= 0 = fmap (const mempty) FL.drain
317+ | otherwise = A.unsafeFreeze <$> Fold step initial done
318+ where
319+ step (Tuple3' rb rh i) a = do
320+ rh1 <- liftIO $ RB.unsafeInsert rb rh a
321+ return $ FL.Partial $ Tuple3' rb rh1 (i + 1)
322+ initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
323+ done (Tuple3' rb rh i) = do
324+ arr <- liftIO $ MA.newArray n
325+ foldFunc i rh snoc' arr rb
326+ snoc' b a = liftIO $ MA.unsafeSnoc b a
327+ foldFunc i
328+ | i < n = RB.unsafeFoldRingM
329+ | otherwise = RB.unsafeFoldRingFullM
330+
304331-------------------------------------------------------------------------------
305332-- Random Access
306333-------------------------------------------------------------------------------
src/Streamly/Internal/Data/Stream/StreamD.hs−24
@@ -299,8 +299,6 @@ module Streamly.Internal.Data.Stream.StreamD
299299 , mkParallel
300300 , mkParallelD
301301 , newCallbackStream
302−
303− , lastN
304302 )
305303where
306304
@@ -337,13 +335,11 @@ import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
337335import Streamly.Internal.Data.Time.Units
338336 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
339337import Streamly.Internal.Data.Unfold.Types (Unfold(..))
340−import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
341338import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
342339
343340import qualified Streamly.Internal.Data.IORef.Prim as Prim
344341import qualified Streamly.Internal.Data.Pipe.Types as Pipe
345342import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
346−import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
347343import qualified Streamly.Internal.Data.Fold as FL
348344import qualified Streamly.Memory.Ring as RB
349345import qualified Streamly.Internal.Data.Stream.StreamK as K
@@ -4819,26 +4815,6 @@ tapAsync f (Stream step1 state1) = Stream step TapInit
48194815 Skip s -> Skip (TapDone s)
48204816 Stop -> Stop
48214817
4822−-- XXX Exported from Array again as this fold is specific to Array
4823−-- | Take last 'n' elements from the stream and discard the rest.
4824−{-# INLINE lastN #-}
4825−lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4826−lastN n
4827− | n <= 0 = fmap (const mempty) FL.drain
4828− | otherwise = A.unsafeFreeze <$> Fold step initial done
4829− where
4830− step (Tuple3' rb rh i) a = do
4831− rh1 <- liftIO $ RB.unsafeInsert rb rh a
4832− return $ FL.Partial $ Tuple3' rb rh1 (i + 1)
4833− initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4834− done (Tuple3' rb rh i) = do
4835− arr <- liftIO $ MA.newArray n
4836− foldFunc i rh snoc' arr rb
4837− snoc' b a = liftIO $ MA.unsafeSnoc b a
4838− foldFunc i
4839− | i < n = RB.unsafeFoldRingM
4840− | otherwise = RB.unsafeFoldRingFullM
4841−
48424818------------------------------------------------------------------------------
48434819-- Time related
48444820------------------------------------------------------------------------------