@@ -0,0 +1,263 @@
1+module Streamly.Internal.Data.Fold.Window
13+import Data.Bifunctor(bimap)
14+import Data.Function ((&))
15+import Data.Maybe (fromMaybe)
17+import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
18+import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
20+import Prelude hiding (length, sum, minimum, maximum)
22+import qualified Deque.Strict as DQ
23+import qualified Streamly.Data.Fold as Fold
26+-- >>> import Data.Bifunctor(bimap)
27+-- >>> import qualified Streamly.Data.Fold as Fold
28+-- >>> import Prelude hiding (length, sum, minimum, maximum)
30+-------------------------------------------------------------------------------
32+-------------------------------------------------------------------------------
34+-- | Map a function on the incoming as well as outgoing element of a rolling
37+-- >>> lmap f = Fold.lmap (bimap f (f <$>))
40+lmap :: (c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
41+lmap f = Fold.lmap (bimap f (f <$>))
43+-- | Convert a rolling fold to a normal fold using the entire input stream as a
46+-- >>> whole f = Fold.lmap (\x -> (x, Nothing)) f
49+whole :: Fold m (a, Maybe a) b -> Fold m a b
50+whole = Fold.lmap (, Nothing)
52+-------------------------------------------------------------------------------
54+-------------------------------------------------------------------------------
58+-- | The sum of all the elements in a rolling window. The input elements are
59+-- required to be intergal numbers.
61+-- This was written in the hope that it would be a tiny bit faster than 'sum'
62+-- for 'Integral' values. But turns out that 'sum' is 2% faster than this even
63+-- for intergal values!
68+sumInt :: forall m a. (Monad m, Integral a) => Fold m (a, Maybe a) a
69+sumInt = Fold step initial extract
73+ initial = return $ Partial (0 :: a)
80+ Just old -> s + a - old
86+-- | Sum of all the elements in a rolling window:
88+-- Uses Kahan-Babuska-Neumaier style summation for numerical stability of
89+-- floating precision arithmetic.
91+-- /Space/: \(\mathcal{O}(1)\)
93+-- /Time/: \(\mathcal{O}(n)\)
96+sum :: forall m a. (Monad m, Num a) => Fold m (a, Maybe a) a
97+sum = Fold step initial extract
105+ (0 :: a) -- running sum
106+ (0 :: a) -- accumulated rounding error
108+ step (Tuple' total err) (new, mOld) =
111+ -- XXX new may be large and err may be small we may lose it
112+ Nothing -> new - err
113+ -- XXX if (new - old) is large we may lose err
114+ Just old -> (new - old) - err
115+ -- total is large and incr may be small, we may round incr here but
116+ -- we will accumulate the rounding error in err1 in the next step.
117+ total1 = total + incr
118+ -- Accumulate any rounding error in err1
119+ -- XXX In the Nothing case above we may lose err, therefore we
120+ -- should use ((total1 - total) - new) + err here.
121+ -- Or even in the just case if (new - old) is large we may lose
122+ -- err, so we should use ((total1 - total) + (old - new)) + err.
123+ err1 = (total1 - total) - incr
124+ in return $ Partial $ Tuple' total1 err1
126+ extract (Tuple' total _) = return total
128+-- | The number of elements in the rolling window.
130+{-# INLINE length #-}
131+length :: (Monad m, Num b) => Fold m (a, Maybe a) b
132+length = Fold.foldl' step 0
136+ step w (_, Nothing) = w + 1
139+-------------------------------------------------------------------------------
141+-------------------------------------------------------------------------------
143+-- Theoretically, we can approximate minimum in a rolling window by using a
144+-- 'powerMean' with sufficiently large negative power.
146+-- XXX If we need to know the minimum in the window only once in a while then
147+-- we can use linear search when it is extracted and not pay the cost all the
150+-- | The minimum element in a rolling window.
152+-- If you want to compute the minimum of the entire stream Fold.minimum from streamly
153+-- package would be much faster.
155+-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
157+{-# INLINE minimum #-}
158+minimum :: (Monad m, Ord a) => Fold m (a, Maybe a) a
159+minimum = Fold step initial extract
163+ initial = return $ Partial $ Tuple3' (0 :: Int) (0 :: Int)
164+ (mempty :: DQ.Deque (Int, a))
166+ step (Tuple3' i w q) (a, ma) =
169+ return $ Partial $ Tuple3' (i + 1) (w + 1)
170+ (headCheck i q (w + 1) & dqloop (i, a))
172+ return $ Partial $ Tuple3' (i + 1) w
173+ (headCheck i q w & dqloop (i,a))
175+ {-# INLINE headCheck #-}
186+ Nothing -> DQ.snoc ia q
187+ -- XXX This can be improved for the case of `=`
189+ if snd ia <= snd ia'
193+ extract (Tuple3' _ _ q) = return $ snd
194+ $ fromMaybe (0, error "min: Empty stream")
197+-- Theoretically, we can approximate maximum in a rolling window by using a
198+-- 'powerMean' with sufficiently large positive power.
200+-- | The maximum element in a rolling window.
202+-- If you want to compute the maximum of the entire stream Fold.maximum from streamly
203+-- package would be much faster.
205+-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
207+{-# INLINE maximum #-}
208+maximum :: (Monad m, Ord a) => Fold m (a, Maybe a) a
209+maximum = Fold step initial extract
213+ initial = return $ Partial $ Tuple3' (0 :: Int) (0 :: Int)
214+ (mempty :: DQ.Deque (Int, a))
216+ step (Tuple3' i w q) (a, ma) =
219+ return $ Partial $ Tuple3' (i + 1) (w + 1)
220+ (headCheck i q (w + 1) & dqloop (i, a))
222+ return $ Partial $ Tuple3' (i + 1) w
223+ (headCheck i q w & dqloop (i,a))
225+ {-# INLINE headCheck #-}
236+ Nothing -> DQ.snoc ia q
237+ -- XXX This can be improved for the case of `=`
239+ if snd ia >= snd ia'
243+ extract (Tuple3' _ _ q) =
246+ $ fromMaybe (0, error "max: Empty stream")
250+-- | The difference between the maximum and minimum elements of a rolling window.
252+-- >>> range = Fold.teeWith (-) maximum minimum
254+-- If you want to compute the range of the entire stream @Fold.teeWith (-)
255+-- Fold.maximum Fold.min@ from the streamly package would be much faster.
257+-- /Space/: \(\mathcal{O}(n)\) where @n@ is the window size.
259+-- /Time/: \(\mathcal{O}(n*w)\) where \(w\) is the window size.
262+range :: (Monad m, Num a, Ord a) => Fold m (a, Maybe a) a
263+range = Fold.teeWith (-) maximum minimum