{-# LANGUAGE CPP #-}
-- |
-- Module      : Streamly.Internal.Data.Fold.Tee
-- Copyright   : (c) 2020 Composewell Technologies
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- A newtype wrapper over the 'Fold' type providing distributing 'Applicative',
-- 'Semigroup', 'Monoid', 'Num', 'Floating' and 'Fractional' instances.
--
module Streamly.Internal.Data.Fold.Tee
    ( Tee(..)
    , toFold
    )
where

#if !MIN_VERSION_base(4,18,0)
import Control.Applicative (liftA2)
#endif
import Streamly.Internal.Data.Fold.Type (Fold)

import qualified Streamly.Internal.Data.Fold.Type as Fold

#include "DocTestDataFold.hs"

-- | @Tee@ is a newtype wrapper over the 'Fold' type providing distributing
-- 'Applicative', 'Semigroup', 'Monoid', 'Num', 'Floating' and 'Fractional'
-- instances.
--
-- The input received by the composed 'Tee' is replicated and distributed to
-- the constituent folds of the 'Tee'.
--
-- For example, to compute the average of numbers in a stream without going
-- through the stream twice:
--
-- >>> avg = (/) <$> (Tee Fold.sum) <*> (Tee $ fmap fromIntegral Fold.length)
-- >>> Stream.fold (unTee avg) $ Stream.fromList [1.0..100.0]
-- 50.5
--
-- Similarly, the 'Semigroup' and 'Monoid' instances of 'Tee' distribute the
-- input to both the folds and combine the outputs using Monoid or Semigroup
-- instances of the output types:
--
-- >>> import Data.Monoid (Sum(..))
-- >>> t = Tee Fold.one <> Tee Fold.latest
-- >>> Stream.fold (unTee t) (fmap Sum $ Stream.enumerateFromTo 1.0 100.0)
-- Just (Sum {getSum = 101.0})
--
-- The 'Num', 'Floating', and 'Fractional' instances work in the same way.
--
newtype Tee m a b =
    Tee { forall (m :: * -> *) a b. Tee m a b -> Fold m a b
unTee :: Fold m a b }
    deriving ((forall a b. (a -> b) -> Tee m a a -> Tee m a b)
-> (forall a b. a -> Tee m a b -> Tee m a a) -> Functor (Tee m a)
forall a b. a -> Tee m a b -> Tee m a a
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (m :: * -> *) a a b.
Functor m =>
a -> Tee m a b -> Tee m a a
forall (m :: * -> *) a a b.
Functor m =>
(a -> b) -> Tee m a a -> Tee m a b
$cfmap :: forall (m :: * -> *) a a b.
Functor m =>
(a -> b) -> Tee m a a -> Tee m a b
fmap :: forall a b. (a -> b) -> Tee m a a -> Tee m a b
$c<$ :: forall (m :: * -> *) a a b.
Functor m =>
a -> Tee m a b -> Tee m a a
<$ :: forall a b. a -> Tee m a b -> Tee m a a
Functor)

{-# DEPRECATED toFold "Please use 'unTee' instead." #-}
toFold :: Tee m a b -> Fold m a b
toFold :: forall (m :: * -> *) a b. Tee m a b -> Fold m a b
toFold = Tee m a b -> Fold m a b
forall (m :: * -> *) a b. Tee m a b -> Fold m a b
unTee

-- | '<*>' distributes the input to both the argument 'Tee's and combines their
-- outputs using function application.
--
instance Monad m => Applicative (Tee m a) where

    {-# INLINE pure #-}
    pure :: forall a. a -> Tee m a a
pure a
a = Fold m a a -> Tee m a a
forall (m :: * -> *) a b. Fold m a b -> Tee m a b
Tee (a -> Fold m a a
forall (m :: * -> *) b a. Applicative m => b -> Fold m a b
Fold.fromPure a
a)

    {-# INLINE (<*>) #-}
    <*> :: forall a b. Tee m a (a -> b) -> Tee m a a -> Tee m a b
(<*>) Tee m a (a -> b)
a Tee m a a
b = Fold m a b -> Tee m a b
forall (m :: * -> *) a b. Fold m a b -> Tee m a b
Tee (((a -> b) -> a -> b)
-> Fold m a (a -> b) -> Fold m a a -> Fold m a b
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 -> b) -> a -> b
forall a b. (a -> b) -> a -> b
($) (Tee m a (a -> b) -> Fold m a (a -> b)
forall (m :: * -> *) a b. Tee m a b -> Fold m a b
unTee Tee m a (a -> b)
a) (Tee m a a -> Fold m a a
forall (m :: * -> *) a b. Tee m a b -> Fold m a b
unTee Tee m a a
b))

-- | '<>' distributes the input to both the argument 'Tee's and combines their
-- outputs using the 'Semigroup' instance of the output type.
--
instance (Semigroup b, Monad m) => Semigroup (Tee m a b) where
    {-# INLINE (<>) #-}
    <> :: Tee m a b -> Tee m a b -> Tee m a b
(<>) = (b -> b -> b) -> Tee m a b -> Tee m a b -> Tee m a b
forall a b c. (a -> b -> c) -> Tee m a a -> Tee m a b -> Tee m a c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> b -> b
forall a. Semigroup a => a -> a -> a
(<>)

-- | '<>' distributes the input to both the argument 'Tee's and combines their
-- outputs using the 'Monoid' instance of the output type.
--
instance (Monoid b, Monad m) => Monoid (Tee m a b) where
    {-# INLINE mempty #-}
    mempty :: Tee m a b
mempty = b -> Tee m a b
forall a. a -> Tee m a a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
forall a. Monoid a => a
mempty

    {-# INLINE mappend #-}
    mappend :: Tee m a b -> Tee m a b -> Tee m a b
mappend = Tee m a b -> Tee m a b -> Tee m a b
forall a. Semigroup a => a -> a -> a
(<>)

-- | Binary 'Num' operations distribute the input to both the argument 'Tee's
-- and combine their outputs using the 'Num' instance of the output type.
--
instance (Monad m, Num b) => Num (Tee m a b) where
    {-# INLINE fromInteger #-}
    fromInteger :: Integer -> Tee m a b
fromInteger = b -> Tee m a b
forall a. a -> Tee m a a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> Tee m a b) -> (Integer -> b) -> Integer -> Tee m a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> b
forall a. Num a => Integer -> a
fromInteger

    {-# INLINE negate #-}
    negate :: Tee m a b -> Tee m a b
negate = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Num a => a -> a
negate

    {-# INLINE abs #-}
    abs :: Tee m a b -> Tee m a b
abs = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Num a => a -> a
abs

    {-# INLINE signum #-}
    signum :: Tee m a b -> Tee m a b
signum = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Num a => a -> a
signum

    {-# INLINE (+) #-}
    + :: Tee m a b -> Tee m a b -> Tee m a b
(+) = (b -> b -> b) -> Tee m a b -> Tee m a b -> Tee m a b
forall a b c. (a -> b -> c) -> Tee m a a -> Tee m a b -> Tee m a c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> b -> b
forall a. Num a => a -> a -> a
(+)

    {-# INLINE (*) #-}
    * :: Tee m a b -> Tee m a b -> Tee m a b
(*) = (b -> b -> b) -> Tee m a b -> Tee m a b -> Tee m a b
forall a b c. (a -> b -> c) -> Tee m a a -> Tee m a b -> Tee m a c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> b -> b
forall a. Num a => a -> a -> a
(*)

    {-# INLINE (-) #-}
    (-) = (b -> b -> b) -> Tee m a b -> Tee m a b -> Tee m a b
forall a b c. (a -> b -> c) -> Tee m a a -> Tee m a b -> Tee m a c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 (-)

-- | Binary 'Fractional' operations distribute the input to both the argument
-- 'Tee's and combine their outputs using the 'Fractional' instance of the
-- output type.
--
instance (Monad m, Fractional b) => Fractional (Tee m a b) where
    {-# INLINE fromRational #-}
    fromRational :: Rational -> Tee m a b
fromRational = b -> Tee m a b
forall a. a -> Tee m a a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> Tee m a b) -> (Rational -> b) -> Rational -> Tee m a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Rational -> b
forall a. Fractional a => Rational -> a
fromRational

    {-# INLINE recip #-}
    recip :: Tee m a b -> Tee m a b
recip = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Fractional a => a -> a
recip

    {-# INLINE (/) #-}
    / :: Tee m a b -> Tee m a b -> Tee m a b
(/) = (b -> b -> b) -> Tee m a b -> Tee m a b -> Tee m a b
forall a b c. (a -> b -> c) -> Tee m a a -> Tee m a b -> Tee m a c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> b -> b
forall a. Fractional a => a -> a -> a
(/)

-- | Binary 'Floating' operations distribute the input to both the argument
-- 'Tee's and combine their outputs using the 'Floating' instance of the output
-- type.
instance (Monad m, Floating b) => Floating (Tee m a b) where
    {-# INLINE pi #-}
    pi :: Tee m a b
pi = b -> Tee m a b
forall a. a -> Tee m a a
forall (f :: * -> *) a. Applicative f => a -> f a
pure b
forall a. Floating a => a
pi

    {-# INLINE exp #-}
    exp :: Tee m a b -> Tee m a b
exp = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
exp

    {-# INLINE sqrt #-}
    sqrt :: Tee m a b -> Tee m a b
sqrt = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
sqrt

    {-# INLINE log #-}
    log :: Tee m a b -> Tee m a b
log = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
log

    {-# INLINE sin #-}
    sin :: Tee m a b -> Tee m a b
sin = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
sin

    {-# INLINE tan #-}
    tan :: Tee m a b -> Tee m a b
tan = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
tan

    {-# INLINE cos #-}
    cos :: Tee m a b -> Tee m a b
cos = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
cos

    {-# INLINE asin #-}
    asin :: Tee m a b -> Tee m a b
asin = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
asin

    {-# INLINE atan #-}
    atan :: Tee m a b -> Tee m a b
atan = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
atan

    {-# INLINE acos #-}
    acos :: Tee m a b -> Tee m a b
acos = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
acos

    {-# INLINE sinh #-}
    sinh :: Tee m a b -> Tee m a b
sinh = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
sinh

    {-# INLINE tanh #-}
    tanh :: Tee m a b -> Tee m a b
tanh = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
tanh

    {-# INLINE cosh #-}
    cosh :: Tee m a b -> Tee m a b
cosh = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
cosh

    {-# INLINE asinh #-}
    asinh :: Tee m a b -> Tee m a b
asinh = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
asinh

    {-# INLINE atanh #-}
    atanh :: Tee m a b -> Tee m a b
atanh = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
atanh

    {-# INLINE acosh #-}
    acosh :: Tee m a b -> Tee m a b
acosh = (b -> b) -> Tee m a b -> Tee m a b
forall a b. (a -> b) -> Tee m a a -> Tee m a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> b
forall a. Floating a => a -> a
acosh

    {-# INLINE (**) #-}
    ** :: Tee m a b -> Tee m a b -> Tee m a b
(**) = (b -> b -> b) -> Tee m a b -> Tee m a b -> Tee m a b
forall a b c. (a -> b -> c) -> Tee m a a -> Tee m a b -> Tee m a c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> b -> b
forall a. Floating a => a -> a -> a
(**)

    {-# INLINE logBase #-}
    logBase :: Tee m a b -> Tee m a b -> Tee m a b
logBase = (b -> b -> b) -> Tee m a b -> Tee m a b -> Tee m a b
forall a b c. (a -> b -> c) -> Tee m a a -> Tee m a b -> Tee m a c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 b -> b -> b
forall a. Floating a => a -> a -> a
logBase