on master help sign in
implement slidingWindow
08f82abAnurag Hooda
src/Streamly/Internal/Data/Ring/Foreign.hs+48−1
@@ -64,6 +64,8 @@ module Streamly.Internal.Data.Ring.Foreign
64 -- * Fast Byte Comparisons
64 -- * Fast Byte Comparisons
65 , unsafeEqArray
65 , unsafeEqArray
66 , unsafeEqArrayN
66 , unsafeEqArrayN
67+
68+ , slidingWindow
67 ) where
69 ) where
68
70
69#include "inline.hs"
71#include "inline.hs"
@@ -78,7 +80,7 @@ import Foreign.Storable (Storable(..))
78import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes)
80import GHC.ForeignPtr (mallocPlainForeignPtrAlignedBytes)
79import GHC.Ptr (Ptr(..))
81import GHC.Ptr (Ptr(..))
80import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array, memcmp)
82import Streamly.Internal.Data.Array.Foreign.Mut.Type (Array, memcmp)
81−import Streamly.Internal.Data.Fold.Type (Fold(..))
83+import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
82import Streamly.Internal.Data.Stream.Serial (SerialT(..))
84import Streamly.Internal.Data.Stream.Serial (SerialT(..))
83import Streamly.Internal.Data.Unfold.Type (Unfold(..))
85import Streamly.Internal.Data.Unfold.Type (Unfold(..))
84import Streamly.Internal.System.IO (unsafeInlineIO)
86import Streamly.Internal.System.IO (unsafeInlineIO)
@@ -506,3 +508,48 @@ unsafeFoldRingNM count rh f z rb@Ring {..} =
506 if ptr == rh || n == 0
508 if ptr == rh || n == 0
507 then return acc'
509 then return acc'
508 else go (n - 1) acc' ptr
510 else go (n - 1) acc' ptr
511+
512+data Tuple4' a b c d = Tuple4' !a !b !c !d deriving Show
513+
514+-- | @slidingWindow collector@ is an incremental sliding window
515+-- fold that does not require all the intermediate elements in a computation.
516+-- This maintains n elements in the window, when a new element comes it slides
517+-- out the oldest element and the new element along with the old element are
518+-- supplied to the collector fold.
519+--
520+-- The Maybe is for the case when initially the window is filling and
521+-- there is no old element.
522+--
523+{-# INLINE slidingWindow #-}
524+slidingWindow :: forall m a b. (MonadIO m, Storable a)
525+ => Int -> Fold m (a, Maybe a) b -> Fold m a b
526+slidingWindow n (Fold step1 initial1 extract1)= Fold step initial extract
527+
528+ where
529+
530+ initial = do
531+ r <- initial1
532+ (rb, rh) <- liftIO $ new n
533+ return $
534+ case r of
535+ Partial s -> Partial $ Tuple4' rb rh (0 :: Int) s
536+ Done b -> Done b
537+
538+ step (Tuple4' rb rh i st) a
539+ | i < n = do
540+ rh1 <- liftIO $ unsafeInsert rb rh a
541+ r <- step1 st (a, Nothing)
542+ return $
543+ case r of
544+ Partial s -> Partial $ Tuple4' rb rh1 (i + 1) s
545+ Done b -> Done b
546+ | otherwise = do
547+ old <- liftIO $ peek rh
548+ rh1 <- liftIO $ unsafeInsert rb rh a
549+ r <- step1 st (a, Just old)
550+ return $
551+ case r of
552+ Partial s -> Partial $ Tuple4' rb rh1 (i + 1) s
553+ Done b -> Done b
554+
555+ extract (Tuple4' _ _ _ st) = extract1 st