@@ -305,6 +305,10 @@ module Streamly.Internal.Data.Stream.StreamD
305305 , mkParallel
306306 , applyParallel
307307 , foldParallel
308+
309+ -- XXX Move this later
310+ -- XXX Exported from Array again
311+ , lastN
308312 )
309313where
310314
@@ -346,6 +350,7 @@ import Streamly.Internal.Data.Fold.Types (Fold(..))
346350import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
347351import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
348352import Streamly.Internal.Data.Unfold.Types (Unfold(..))
353+import Streamly.Internal.Data.Strict (Tuple3'(..))
349354
350355import Streamly.Internal.Data.Stream.StreamD.Type
351356import Streamly.Internal.Data.SVar
@@ -4013,3 +4018,21 @@ tapAsync f (Stream step1 state1) = Stream step TapInit
40134018 Yield a s -> Yield a (TapDone s)
40144019 Skip s -> Skip (TapDone s)
40154020 Stop -> Stop
4021+
4022+-- XXX Is there room for improvement?
4023+-- | Take last 'n' elements from the stream and discard the rest.
4024+{-# INLINABLE lastN #-}
4025+lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4026+lastN n = Fold step initial done
4027+ where
4028+ step (Tuple3' rb rh i) a = do
4029+ rh1 <- liftIO $ RB.unsafeInsert rb rh a
4030+ return $ Tuple3' rb rh1 (i + 1)
4031+ initial = fmap (\(a, b) -> Tuple3' a b 0) $ liftIO $ RB.new n
4032+ done (Tuple3' rb rh i) = do
4033+ arr <- liftIO $ A.newArray n
4034+ foldFunc i rh snoc' arr rb
4035+ snoc' b a = liftIO $ A.unsafeSnoc b a
4036+ foldFunc i
4037+ | i < n = RB.unsafeFoldRingM
4038+ | otherwise = RB.unsafeFoldRingFullM