module Streamly.Internal.Data.Fold.Window
(
lmap
, cumulative
, rollingMap
, rollingMapM
, length
, sum
, sumInt
, powerSum
, powerSumFrac
, minimum
, maximum
, range
, mean
)
where
import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.Bifunctor(bimap)
import Foreign.Storable (Storable, peek)
import Streamly.Internal.Data.Fold.Type (Fold(..), Step(..))
import Streamly.Internal.Data.Tuple.Strict
(Tuple'(..), Tuple3Fused' (Tuple3Fused'))
import qualified Streamly.Internal.Data.Fold.Type as Fold
import qualified Streamly.Internal.Data.Ring.Unboxed as Ring
import Prelude hiding (length, sum, minimum, maximum)
{-# INLINE lmap #-}
lmap :: (c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
lmap :: (c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
lmap c -> a
f = ((c, Maybe c) -> (a, Maybe a))
-> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
forall a b (m :: * -> *) r. (a -> b) -> Fold m b r -> Fold m a r
Fold.lmap ((c -> a) -> (Maybe c -> Maybe a) -> (c, Maybe c) -> (a, Maybe a)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap c -> a
f (c -> a
f (c -> a) -> Maybe c -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>))
{-# INLINE cumulative #-}
cumulative :: Fold m (a, Maybe a) b -> Fold m a b
cumulative :: Fold m (a, Maybe a) b -> Fold m a b
cumulative = (a -> (a, Maybe a)) -> Fold m (a, Maybe a) b -> Fold m a b
forall a b (m :: * -> *) r. (a -> b) -> Fold m b r -> Fold m a r
Fold.lmap (, Maybe a
forall a. Maybe a
Nothing)
{-# INLINE rollingMapM #-}
rollingMapM :: Monad m =>
(Maybe a -> a -> m (Maybe b)) -> Fold m (a, Maybe a) (Maybe b)
rollingMapM :: (Maybe a -> a -> m (Maybe b)) -> Fold m (a, Maybe a) (Maybe b)
rollingMapM Maybe a -> a -> m (Maybe b)
f = (Maybe b -> (a, Maybe a) -> m (Maybe b))
-> m (Maybe b) -> Fold m (a, Maybe a) (Maybe b)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Fold m a b
Fold.foldlM' Maybe b -> (a, Maybe a) -> m (Maybe b)
forall p. p -> (a, Maybe a) -> m (Maybe b)
f1 m (Maybe b)
forall a. m (Maybe a)
initial
where
initial :: m (Maybe a)
initial = Maybe a -> m (Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe a
forall a. Maybe a
Nothing
f1 :: p -> (a, Maybe a) -> m (Maybe b)
f1 p
_ (a
a, Maybe a
ma) = Maybe a -> a -> m (Maybe b)
f Maybe a
ma a
a
{-# INLINE rollingMap #-}
rollingMap :: Monad m =>
(Maybe a -> a -> Maybe b) -> Fold m (a, Maybe a) (Maybe b)
rollingMap :: (Maybe a -> a -> Maybe b) -> Fold m (a, Maybe a) (Maybe b)
rollingMap Maybe a -> a -> Maybe b
f = (Maybe b -> (a, Maybe a) -> Maybe b)
-> Maybe b -> Fold m (a, Maybe a) (Maybe b)
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Fold m a b
Fold.foldl' Maybe b -> (a, Maybe a) -> Maybe b
forall p. p -> (a, Maybe a) -> Maybe b
f1 Maybe b
forall a. Maybe a
initial
where
initial :: Maybe a
initial = Maybe a
forall a. Maybe a
Nothing
f1 :: p -> (a, Maybe a) -> Maybe b
f1 p
_ (a
a, Maybe a
ma) = Maybe a -> a -> Maybe b
f Maybe a
ma a
a
{-# INLINE sumInt #-}
sumInt :: forall m a. (Monad m, Integral a) => Fold m (a, Maybe a) a
sumInt :: Fold m (a, Maybe a) a
sumInt = (a -> (a, Maybe a) -> m (Step a a))
-> m (Step a a) -> (a -> m a) -> Fold m (a, Maybe a) a
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold a -> (a, Maybe a) -> m (Step a a)
forall (m :: * -> *) a b.
(Monad m, Num a) =>
a -> (a, Maybe a) -> m (Step a b)
step m (Step a a)
forall b. m (Step a b)
initial a -> m a
forall a. a -> m a
extract
where
initial :: m (Step a b)
initial = Step a b -> m (Step a b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step a b -> m (Step a b)) -> Step a b -> m (Step a b)
forall a b. (a -> b) -> a -> b
$ a -> Step a b
forall s b. s -> Step s b
Partial (a
0 :: a)
step :: a -> (a, Maybe a) -> m (Step a b)
step a
s (a
a, Maybe a
ma) =
Step a b -> m (Step a b)
forall (m :: * -> *) a. Monad m => a -> m a
return
(Step a b -> m (Step a b)) -> Step a b -> m (Step a b)
forall a b. (a -> b) -> a -> b
$ a -> Step a b
forall s b. s -> Step s b
Partial
(a -> Step a b) -> a -> Step a b
forall a b. (a -> b) -> a -> b
$ case Maybe a
ma of
Maybe a
Nothing -> a
s a -> a -> a
forall a. Num a => a -> a -> a
+ a
a
Just a
old -> a
s a -> a -> a
forall a. Num a => a -> a -> a
+ a
a a -> a -> a
forall a. Num a => a -> a -> a
- a
old
extract :: a -> m a
extract = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return
{-# INLINE sum #-}
sum :: forall m a. (Monad m, Num a) => Fold m (a, Maybe a) a
sum :: Fold m (a, Maybe a) a
sum = (Tuple' a a -> (a, Maybe a) -> m (Step (Tuple' a a) a))
-> m (Step (Tuple' a a) a)
-> (Tuple' a a -> m a)
-> Fold m (a, Maybe a) a
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold Tuple' a a -> (a, Maybe a) -> m (Step (Tuple' a a) a)
forall (m :: * -> *) b b.
(Monad m, Num b) =>
Tuple' b b -> (b, Maybe b) -> m (Step (Tuple' b b) b)
step m (Step (Tuple' a a) a)
forall b. m (Step (Tuple' a a) b)
initial Tuple' a a -> m a
forall (m :: * -> *) a b. Monad m => Tuple' a b -> m a
extract
where
initial :: m (Step (Tuple' a a) b)
initial =
Step (Tuple' a a) b -> m (Step (Tuple' a a) b)
forall (m :: * -> *) a. Monad m => a -> m a
return
(Step (Tuple' a a) b -> m (Step (Tuple' a a) b))
-> Step (Tuple' a a) b -> m (Step (Tuple' a a) b)
forall a b. (a -> b) -> a -> b
$ Tuple' a a -> Step (Tuple' a a) b
forall s b. s -> Step s b
Partial
(Tuple' a a -> Step (Tuple' a a) b)
-> Tuple' a a -> Step (Tuple' a a) b
forall a b. (a -> b) -> a -> b
$ a -> a -> Tuple' a a
forall a b. a -> b -> Tuple' a b
Tuple'
(a
0 :: a)
(a
0 :: a)
step :: Tuple' b b -> (b, Maybe b) -> m (Step (Tuple' b b) b)
step (Tuple' b
total b
err) (b
new, Maybe b
mOld) =
let incr :: b
incr =
case Maybe b
mOld of
Maybe b
Nothing -> b
new b -> b -> b
forall a. Num a => a -> a -> a
- b
err
Just old -> (b
new b -> b -> b
forall a. Num a => a -> a -> a
- b
old) b -> b -> b
forall a. Num a => a -> a -> a
- b
err
total1 :: b
total1 = b
total b -> b -> b
forall a. Num a => a -> a -> a
+ b
incr
err1 :: b
err1 = (b
total1 b -> b -> b
forall a. Num a => a -> a -> a
- b
total) b -> b -> b
forall a. Num a => a -> a -> a
- b
incr
in Step (Tuple' b b) b -> m (Step (Tuple' b b) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' b b) b -> m (Step (Tuple' b b) b))
-> Step (Tuple' b b) b -> m (Step (Tuple' b b) b)
forall a b. (a -> b) -> a -> b
$ Tuple' b b -> Step (Tuple' b b) b
forall s b. s -> Step s b
Partial (Tuple' b b -> Step (Tuple' b b) b)
-> Tuple' b b -> Step (Tuple' b b) b
forall a b. (a -> b) -> a -> b
$ b -> b -> Tuple' b b
forall a b. a -> b -> Tuple' a b
Tuple' b
total1 b
err1
extract :: Tuple' a b -> m a
extract (Tuple' a
total b
_) = a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
total
{-# INLINE length #-}
length :: (Monad m, Num b) => Fold m (a, Maybe a) b
length :: Fold m (a, Maybe a) b
length = (b -> (a, Maybe a) -> b) -> b -> Fold m (a, Maybe a) b
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Fold m a b
Fold.foldl' b -> (a, Maybe a) -> b
forall p a a. Num p => p -> (a, Maybe a) -> p
step b
0
where
step :: p -> (a, Maybe a) -> p
step p
w (a
_, Maybe a
Nothing) = p
w p -> p -> p
forall a. Num a => a -> a -> a
+ p
1
step p
w (a, Maybe a)
_ = p
w
{-# INLINE powerSum #-}
powerSum :: (Monad m, Num a) => Int -> Fold m (a, Maybe a) a
powerSum :: Int -> Fold m (a, Maybe a) a
powerSum Int
k = (a -> a) -> Fold m (a, Maybe a) a -> Fold m (a, Maybe a) a
forall c a (m :: * -> *) b.
(c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
lmap (a -> Int -> a
forall a b. (Num a, Integral b) => a -> b -> a
^ Int
k) Fold m (a, Maybe a) a
forall (m :: * -> *) a. (Monad m, Num a) => Fold m (a, Maybe a) a
sum
{-# INLINE powerSumFrac #-}
powerSumFrac :: (Monad m, Floating a) => a -> Fold m (a, Maybe a) a
powerSumFrac :: a -> Fold m (a, Maybe a) a
powerSumFrac a
p = (a -> a) -> Fold m (a, Maybe a) a -> Fold m (a, Maybe a) a
forall c a (m :: * -> *) b.
(c -> a) -> Fold m (a, Maybe a) b -> Fold m (c, Maybe c) b
lmap (a -> a -> a
forall a. Floating a => a -> a -> a
** a
p) Fold m (a, Maybe a) a
forall (m :: * -> *) a. (Monad m, Num a) => Fold m (a, Maybe a) a
sum
{-# INLINE range #-}
range :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe (a, a))
range :: Int -> Fold m a (Maybe (a, a))
range Int
n = (Tuple3Fused' (Ring a) (Ptr a) Int
-> a
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) (Maybe (a, a))))
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) (Maybe (a, a)))
-> (Tuple3Fused' (Ring a) (Ptr a) Int -> m (Maybe (a, a)))
-> Fold m a (Maybe (a, a))
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold Tuple3Fused' (Ring a) (Ptr a) Int
-> a -> m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) (Maybe (a, a)))
forall (m :: * -> *) a c b.
(MonadIO m, Storable a, Num c) =>
Tuple3Fused' (Ring a) (Ptr a) c
-> a -> m (Step (Tuple3Fused' (Ring a) (Ptr a) c) b)
step m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) (Maybe (a, a)))
forall b. m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
initial Tuple3Fused' (Ring a) (Ptr a) Int -> m (Maybe (a, a))
forall (m :: * -> *) b.
(MonadIO m, Storable b, Ord b) =>
Tuple3Fused' (Ring b) (Ptr b) Int -> m (Maybe (b, b))
extract
where
initial :: m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
initial =
if Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0
then [Char] -> m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
forall a. HasCallStack => [Char] -> a
error [Char]
"range: window size must be > 0"
else
let f :: (a, b) -> Step (Tuple3Fused' a b Int) b
f (a
a, b
b) = Tuple3Fused' a b Int -> Step (Tuple3Fused' a b Int) b
forall s b. s -> Step s b
Partial (Tuple3Fused' a b Int -> Step (Tuple3Fused' a b Int) b)
-> Tuple3Fused' a b Int -> Step (Tuple3Fused' a b Int) b
forall a b. (a -> b) -> a -> b
$ a -> b -> Int -> Tuple3Fused' a b Int
forall a b c. a -> b -> c -> Tuple3Fused' a b c
Tuple3Fused' a
a b
b (Int
0 :: Int)
in ((Ring a, Ptr a) -> Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
-> m (Ring a, Ptr a)
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Ring a, Ptr a) -> Step (Tuple3Fused' (Ring a) (Ptr a) Int) b
forall a b b. (a, b) -> Step (Tuple3Fused' a b Int) b
f (m (Ring a, Ptr a)
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b))
-> m (Ring a, Ptr a)
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) Int) b)
forall a b. (a -> b) -> a -> b
$ IO (Ring a, Ptr a) -> m (Ring a, Ptr a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Ring a, Ptr a) -> m (Ring a, Ptr a))
-> IO (Ring a, Ptr a) -> m (Ring a, Ptr a)
forall a b. (a -> b) -> a -> b
$ Int -> IO (Ring a, Ptr a)
forall a. Storable a => Int -> IO (Ring a, Ptr a)
Ring.new Int
n
step :: Tuple3Fused' (Ring a) (Ptr a) c
-> a -> m (Step (Tuple3Fused' (Ring a) (Ptr a) c) b)
step (Tuple3Fused' Ring a
rb Ptr a
rh c
i) a
a = do
Ptr a
rh1 <- IO (Ptr a) -> m (Ptr a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Ptr a) -> m (Ptr a)) -> IO (Ptr a) -> m (Ptr a)
forall a b. (a -> b) -> a -> b
$ Ring a -> Ptr a -> a -> IO (Ptr a)
forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
Ring.unsafeInsert Ring a
rb Ptr a
rh a
a
Step (Tuple3Fused' (Ring a) (Ptr a) c) b
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) c) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple3Fused' (Ring a) (Ptr a) c) b
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) c) b))
-> Step (Tuple3Fused' (Ring a) (Ptr a) c) b
-> m (Step (Tuple3Fused' (Ring a) (Ptr a) c) b)
forall a b. (a -> b) -> a -> b
$ Tuple3Fused' (Ring a) (Ptr a) c
-> Step (Tuple3Fused' (Ring a) (Ptr a) c) b
forall s b. s -> Step s b
Partial (Tuple3Fused' (Ring a) (Ptr a) c
-> Step (Tuple3Fused' (Ring a) (Ptr a) c) b)
-> Tuple3Fused' (Ring a) (Ptr a) c
-> Step (Tuple3Fused' (Ring a) (Ptr a) c) b
forall a b. (a -> b) -> a -> b
$ Ring a -> Ptr a -> c -> Tuple3Fused' (Ring a) (Ptr a) c
forall a b c. a -> b -> c -> Tuple3Fused' a b c
Tuple3Fused' Ring a
rb Ptr a
rh1 (c
i c -> c -> c
forall a. Num a => a -> a -> a
+ c
1)
foldFunc :: Int -> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
foldFunc Int
i
| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n = Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
Ring.unsafeFoldRingM
| Bool
otherwise = Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
Ring.unsafeFoldRingFullM
extract :: Tuple3Fused' (Ring b) (Ptr b) Int -> m (Maybe (b, b))
extract (Tuple3Fused' Ring b
rb Ptr b
rh Int
i) =
if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
then Maybe (b, b) -> m (Maybe (b, b))
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe (b, b)
forall a. Maybe a
Nothing
else do
b
x <- IO b -> m b
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO b -> m b) -> IO b -> m b
forall a b. (a -> b) -> a -> b
$ Ptr b -> IO b
forall a. Storable a => Ptr a -> IO a
peek Ptr b
rh
let accum :: (b, b) -> b -> m (b, b)
accum (b
mn, b
mx) b
a = (b, b) -> m (b, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> b -> b
forall a. Ord a => a -> a -> a
min b
mn b
a, b -> b -> b
forall a. Ord a => a -> a -> a
max b
mx b
a)
((b, b) -> Maybe (b, b)) -> m (b, b) -> m (Maybe (b, b))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (b, b) -> Maybe (b, b)
forall a. a -> Maybe a
Just (m (b, b) -> m (Maybe (b, b))) -> m (b, b) -> m (Maybe (b, b))
forall a b. (a -> b) -> a -> b
$ Int
-> Ptr b
-> ((b, b) -> b -> m (b, b))
-> (b, b)
-> Ring b
-> m (b, b)
forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Int -> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
foldFunc Int
i Ptr b
rh (b, b) -> b -> m (b, b)
forall (m :: * -> *) b. (Monad m, Ord b) => (b, b) -> b -> m (b, b)
accum (b
x, b
x) Ring b
rb
{-# INLINE minimum #-}
minimum :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe a)
minimum :: Int -> Fold m a (Maybe a)
minimum Int
n = (Maybe (a, a) -> Maybe a)
-> Fold m a (Maybe (a, a)) -> Fold m a (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((a, a) -> a) -> Maybe (a, a) -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, a) -> a
forall a b. (a, b) -> a
fst) (Fold m a (Maybe (a, a)) -> Fold m a (Maybe a))
-> Fold m a (Maybe (a, a)) -> Fold m a (Maybe a)
forall a b. (a -> b) -> a -> b
$ Int -> Fold m a (Maybe (a, a))
forall (m :: * -> *) a.
(MonadIO m, Storable a, Ord a) =>
Int -> Fold m a (Maybe (a, a))
range Int
n
{-# INLINE maximum #-}
maximum :: (MonadIO m, Storable a, Ord a) => Int -> Fold m a (Maybe a)
maximum :: Int -> Fold m a (Maybe a)
maximum Int
n = (Maybe (a, a) -> Maybe a)
-> Fold m a (Maybe (a, a)) -> Fold m a (Maybe a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((a, a) -> a) -> Maybe (a, a) -> Maybe a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, a) -> a
forall a b. (a, b) -> b
snd) (Fold m a (Maybe (a, a)) -> Fold m a (Maybe a))
-> Fold m a (Maybe (a, a)) -> Fold m a (Maybe a)
forall a b. (a -> b) -> a -> b
$ Int -> Fold m a (Maybe (a, a))
forall (m :: * -> *) a.
(MonadIO m, Storable a, Ord a) =>
Int -> Fold m a (Maybe (a, a))
range Int
n
{-# INLINE mean #-}
mean :: forall m a. (Monad m, Fractional a) => Fold m (a, Maybe a) a
mean :: Fold m (a, Maybe a) a
mean = (a -> a -> a)
-> Fold m (a, Maybe a) a
-> Fold m (a, Maybe a) a
-> Fold m (a, Maybe a) a
forall (m :: * -> *) a b c x.
Monad m =>
(a -> b -> c) -> Fold m x a -> Fold m x b -> Fold m x c
Fold.teeWith a -> a -> a
forall a. Fractional a => a -> a -> a
(/) Fold m (a, Maybe a) a
forall (m :: * -> *) a. (Monad m, Num a) => Fold m (a, Maybe a) a
sum Fold m (a, Maybe a) a
forall (m :: * -> *) b a. (Monad m, Num b) => Fold m (a, Maybe a) b
length