{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE ConstraintKinds           #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE PatternSynonyms           #-}
{-# LANGUAGE RecordWildCards           #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE ViewPatterns              #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE MagicHash                 #-}

#if __GLASGOW_HASKELL__ >= 801
{-# LANGUAGE TypeApplications          #-}
#endif

#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Stream.StreamD
-- Copyright   : (c) 2018 Harendra Kumar
--               (c) Roman Leshchinskiy 2008-2010
--               (c) The University of Glasgow, 2009
--
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- Direct style re-implementation of CPS style stream in StreamK module.  The
-- symbol or suffix 'D' in this module denotes the "Direct" style.  GHC is able
-- to INLINE and fuse direct style better, providing better performance than
-- CPS implementation.
--
-- @
-- import qualified Streamly.Internal.Data.Stream.StreamD as D
-- @

-- Some of the functions in this file have been adapted from the vector
-- library,  https://hackage.haskell.org/package/vector.

module Streamly.Internal.Data.Stream.StreamD
    (
    -- * The stream type
      Step (..)

#if __GLASGOW_HASKELL__ >= 800
    , Stream (Stream, UnStream)
#else
    , Stream (UnStream)
    , pattern Stream
#endif

    -- * Construction
    , nil
    , nilM
    , cons

    -- * Deconstruction
    , uncons

    -- * Generation
    -- ** Unfolds
    , unfoldr
    , unfoldrM
    , unfold

    -- ** Specialized Generation
    -- | Generate a monadic stream from a seed.
    , repeat
    , repeatM
    , replicate
    , replicateM
    , fromIndices
    , fromIndicesM
    , generate
    , generateM
    , iterate
    , iterateM

    -- ** Enumerations
    , enumerateFromStepIntegral
    , enumerateFromIntegral
    , enumerateFromThenIntegral
    , enumerateFromToIntegral
    , enumerateFromThenToIntegral

    , enumerateFromStepNum
    , numFrom
    , numFromThen
    , enumerateFromToFractional
    , enumerateFromThenToFractional

    -- ** Time
    , currentTime

    -- ** Conversions
    -- | Transform an input structure into a stream.
    -- | Direct style stream does not support @fromFoldable@.
    , yield
    , yieldM
    , fromList
    , fromListM
    , fromStreamK
    , fromStreamD
    , fromPrimVar
    , fromSVar

    -- * Elimination
    -- ** General Folds
    , foldrS
    , foldrT
    , foldrM
    , foldrMx
    , foldr
    , foldr1

    , foldl'
    , foldlM'
    , foldlS
    , foldlT
    , reverse
    , reverse'

    , foldlx'
    , foldlMx'
    , runFold

    , parselMx'
    , splitParse

    -- ** Specialized Folds
    , tap
    , tapOffsetEvery
    , tapAsync
    , tapRate
    , pollCounts
    , drain
    , null
    , head
    , headElse
    , tail
    , last
    , elem
    , notElem
    , all
    , any
    , maximum
    , maximumBy
    , minimum
    , minimumBy
    , findIndices
    , lookup
    , findM
    , find
    , (!!)
    , toSVarParallel

    -- ** Flattening nested streams
    , concatMapM
    , concatMap
    , ConcatMapUState (..)
    , concatMapU
    , ConcatUnfoldInterleaveState (..)
    , concatUnfoldInterleave
    , concatUnfoldRoundrobin
    , AppendState(..)
    , append
    , InterleaveState(..)
    , interleave
    , interleaveMin
    , interleaveSuffix
    , interleaveInfix
    , roundRobin -- interleaveFair?/ParallelFair
    , gintercalateSuffix
    , interposeSuffix
    , gintercalate
    , interpose

    -- ** Grouping
    , groupsOf
    , groupsOf2
    , groupsBy
    , groupsRollingBy

    -- ** Splitting
    , splitBy
    , splitSuffixBy
    , wordsBy
    , splitSuffixBy'

    , splitOn
    , splitSuffixOn

    , splitInnerBy
    , splitInnerBySuffix

    -- ** Substreams
    , isPrefixOf
    , isSubsequenceOf
    , stripPrefix

    -- ** Map and Fold
    , mapM_

    -- ** Conversions
    -- | Transform a stream into another type.
    , toList
    , toListRev
    , toStreamK
    , toStreamD

    , hoist
    , generally

    , liftInner
    , runReaderT
    , evalStateT
    , runStateT

    -- * Transformation
    , transform

    -- ** By folding (scans)
    , scanlM'
    , scanl'
    , scanlM
    , scanl
    , scanl1M'
    , scanl1'
    , scanl1M
    , scanl1

    , prescanl'
    , prescanlM'

    , postscanl
    , postscanlM
    , postscanl'
    , postscanlM'

    , postscanlx'
    , postscanlMx'
    , scanlMx'
    , scanlx'

    -- * Filtering
    , filter
    , filterM
    , uniq
    , take
    , takeByTime
    , takeWhile
    , takeWhileM
    , drop
    , dropByTime
    , dropWhile
    , dropWhileM

    -- * Mapping
    , map
    , mapM
    , sequence
    , rollingMap
    , rollingMapM

    -- * Inserting
    , intersperseM
    , intersperse
    , intersperseSuffix
    , intersperseSuffixBySpan
    , insertBy

    -- * Deleting
    , deleteBy

    -- ** Map and Filter
    , mapMaybe
    , mapMaybeM

    -- * Zipping
    , indexed
    , indexedR
    , zipWith
    , zipWithM

    -- * Comparisons
    , eqBy
    , cmpBy

    -- * Merging
    , mergeBy
    , mergeByM

    -- * Transformation comprehensions
    , the

    -- * Exceptions
    , newFinalizedIORef
    , runIORefFinalizer
    , clearIORefFinalizer
    , gbracket
    , before
    , after
    , afterIO
    , bracket
    , bracketIO
    , onException
    , finally
    , finallyIO
    , handle

    -- * Concurrent Application
    , mkParallel
    , mkParallelD
    , newCallbackStream

    , lastN
    )
where

import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
import Control.Exception
       (assert, Exception, SomeException, AsyncException, fromException, mask_)
import Control.Monad (void, when, forever)
import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Reader (ReaderT)
import Control.Monad.State.Strict (StateT)
import Control.Monad.Trans (MonadTrans(lift))
import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
import Data.Bits (shiftR, shiftL, (.|.), (.&.))
import Data.Functor.Identity (Identity(..))
import Data.Int (Int64)
import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
import Data.Maybe (fromJust, isJust, isNothing)
import Data.Word (Word32)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable(..))
import GHC.Types (SPEC(..))
import System.Mem (performMajorGC)
import Prelude
       hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
               takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
               notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
               (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
               reverse, iterate, splitAt)

import qualified Control.Monad.Catch as MC
import qualified Control.Monad.Reader as Reader
import qualified Control.Monad.State.Strict as State
import qualified Prelude

import Fusion.Plugin.Types (Fuse(..))
import Streamly.Internal.Mutable.Prim.Var
       (Prim, Var, readVar, newVar, modifyVar')
import Streamly.Internal.Data.Time.Units
       (TimeUnit64, toRelTime64, diffAbsTime64)

import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
import Streamly.Internal.Memory.Array.Types (Array(..))
import Streamly.Internal.Data.Fold.Types (Fold(..))
import Streamly.Internal.Data.Parser.Types (Parser(..), ParseError(..))
import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
import Streamly.Internal.Data.Time.Units
       (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
import Streamly.Internal.Data.Unfold.Types (Unfold(..))
import Streamly.Internal.Data.Strict (Tuple3'(..))

import Streamly.Internal.Data.Stream.StreamD.Type
import Streamly.Internal.Data.SVar
import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)

import qualified Streamly.Internal.Data.Pipe.Types as Pipe
import qualified Streamly.Internal.Memory.Array.Types as A
import qualified Streamly.Internal.Data.Fold as FL
import qualified Streamly.Memory.Ring as RB
import qualified Streamly.Internal.Data.Stream.StreamK as K
import qualified Streamly.Internal.Data.Parser.Types as PR

------------------------------------------------------------------------------
-- Construction
------------------------------------------------------------------------------

-- | An empty 'Stream'.
{-# INLINE_NORMAL nil #-}
nil :: Monad m => Stream m a
nil :: forall (m :: * -> *) a. Monad m => Stream m a
nil = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop) ()

-- | An empty 'Stream' with a side effect.
{-# INLINE_NORMAL nilM #-}
nilM :: Monad m => m b -> Stream m a
nilM :: forall (m :: * -> *) b a. Monad m => m b -> Stream m a
nilM m b
m = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> m b
m forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop) ()

{-# INLINE_NORMAL consM #-}
consM :: Monad m => m a -> Stream m a -> Stream m a
consM :: forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
consM m a
m (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 forall a. Maybe a
Nothing
    where
    {-# INLINE_LATE step1 #-}
    step1 :: State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 State Stream m a
_ Maybe s
Nothing   = m a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just s
state)
    step1 State Stream m a
gst (Just s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
          case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall a. a -> Maybe a
Just s
s)
            Skip  s
s   -> forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
s)
            Step s a
Stop      -> forall s a. Step s a
Stop

-- XXX implement in terms of consM?
-- cons x = consM (return x)
--
-- | Can fuse but has O(n^2) complexity.
{-# INLINE_NORMAL cons #-}
cons :: Monad m => a -> Stream m a -> Stream m a
cons :: forall (m :: * -> *) a. Monad m => a -> Stream m a -> Stream m a
cons a
x (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 forall a. Maybe a
Nothing
    where
    {-# INLINE_LATE step1 #-}
    step1 :: State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 State Stream m a
_ Maybe s
Nothing   = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just s
state)
    step1 State Stream m a
gst (Just s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
          case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall a. a -> Maybe a
Just s
s)
            Skip  s
s   -> forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
s)
            Step s a
Stop      -> forall s a. Step s a
Stop

-------------------------------------------------------------------------------
-- Deconstruction
-------------------------------------------------------------------------------

-- Does not fuse, has the same performance as the StreamK version.
{-# INLINE_NORMAL uncons #-}
uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
uncons :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> m (Maybe (a, Stream m a))
uncons (UnStream State Stream m a -> s -> m (Step s a)
step s
state) = s -> m (Maybe (a, Stream m a))
go s
state
  where
    go :: s -> m (Maybe (a, Stream m a))
go s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (a
x, forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step s
s)
            Skip  s
s   -> s -> m (Maybe (a, Stream m a))
go s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

------------------------------------------------------------------------------
-- Generation by unfold
------------------------------------------------------------------------------

{-# INLINE_NORMAL unfoldrM #-}
unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
unfoldrM :: forall (m :: * -> *) s a.
Monad m =>
(s -> m (Maybe (a, s))) -> s -> Stream m a
unfoldrM s -> m (Maybe (a, s))
next s
state = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> s -> m (Step s a)
step s
state
  where
    {-# INLINE_LATE step #-}
    step :: p -> s -> m (Step s a)
step p
_ s
st = do
        Maybe (a, s)
r <- s -> m (Maybe (a, s))
next s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (a, s)
r of
            Just (a
x, s
s) -> forall s a. a -> s -> Step s a
Yield a
x s
s
            Maybe (a, s)
Nothing     -> forall s a. Step s a
Stop

{-# INLINE_LATE unfoldr #-}
unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
unfoldr :: forall (m :: * -> *) s a.
Monad m =>
(s -> Maybe (a, s)) -> s -> Stream m a
unfoldr s -> Maybe (a, s)
f = forall (m :: * -> *) s a.
Monad m =>
(s -> m (Maybe (a, s))) -> s -> Stream m a
unfoldrM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. s -> Maybe (a, s)
f)

-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
--
{-# INLINE_NORMAL unfold #-}
unfold :: Monad m => Unfold m a b -> a -> Stream m b
unfold :: forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> a -> Stream m b
unfold (Unfold s -> m (Step s b)
ustep a -> m s
inject) a
seed = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> Maybe s -> m (Step (Maybe s) b)
step forall a. Maybe a
Nothing
  where
    {-# INLINE_LATE step #-}
    step :: p -> Maybe s -> m (Step (Maybe s) b)
step p
_ Maybe s
Nothing = a -> m s
inject a
seed forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. s -> Step s a
Skip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a
Just
    step p
_ (Just s
st) = do
        Step s b
r <- s -> m (Step s b)
ustep s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall a. a -> Maybe a
Just s
s)
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
s)
            Step s b
Stop      -> forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Specialized Generation
------------------------------------------------------------------------------

{-# INLINE_NORMAL repeatM #-}
repeatM :: Monad m => m a -> Stream m a
repeatM :: forall (m :: * -> *) a. Monad m => m a -> Stream m a
repeatM m a
x = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> m a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r ()) ()

{-# INLINE_NORMAL repeat #-}
repeat :: Monad m => a -> Stream m a
repeat :: forall (m :: * -> *) a. Monad m => a -> Stream m a
repeat a
x = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x ()) ()

{-# INLINE_NORMAL iterateM #-}
iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
iterateM :: forall (m :: * -> *) a. Monad m => (a -> m a) -> m a -> Stream m a
iterateM a -> m a
step = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ m a
st -> m a
st forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (a -> m a
step a
x))

{-# INLINE_NORMAL iterate #-}
iterate :: Monad m => (a -> a) -> a -> Stream m a
iterate :: forall (m :: * -> *) a. Monad m => (a -> a) -> a -> Stream m a
iterate a -> a
step a
st = forall (m :: * -> *) a. Monad m => (a -> m a) -> m a -> Stream m a
iterateM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> a
step) (forall (m :: * -> *) a. Monad m => a -> m a
return a
st)

{-# INLINE_NORMAL replicateM #-}
replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
replicateM :: forall (m :: * -> *) a. Monad m => Int -> m a -> Stream m a
replicateM Int
n m a
p = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> Int -> m (Step Int a)
step Int
n
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ (Int
i :: Int)
      | Int
i forall a. Ord a => a -> a -> Bool
<= Int
0    = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
      | Bool
otherwise = do
          a
x <- m a
p
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (Int
i forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE_NORMAL replicate #-}
replicate :: Monad m => Int -> a -> Stream m a
replicate :: forall (m :: * -> *) a. Monad m => Int -> a -> Stream m a
replicate Int
n a
x = forall (m :: * -> *) a. Monad m => Int -> m a -> Stream m a
replicateM Int
n (forall (m :: * -> *) a. Monad m => a -> m a
return a
x)

-- This would not work properly for floats, therefore we put an Integral
-- constraint.
-- | Can be used to enumerate unbounded integrals. This does not check for
-- overflow or underflow for bounded integrals.
{-# INLINE_NORMAL enumerateFromStepIntegral #-}
enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
enumerateFromStepIntegral :: forall a (m :: * -> *).
(Integral a, Monad m) =>
a -> a -> Stream m a
enumerateFromStepIntegral a
from a
stride =
    a
from seq :: forall a b. a -> b -> b
`seq` a
stride seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p}. Monad m => p -> a -> m (Step a a)
step a
from
    where
        {-# INLINE_LATE step #-}
        step :: p -> a -> m (Step a a)
step p
_ !a
x = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x forall a b. (a -> b) -> a -> b
$! (a
x forall a. Num a => a -> a -> a
+ a
stride)

-- We are assuming that "to" is constrained by the type to be within
-- max/min bounds.
{-# INLINE enumerateFromToIntegral #-}
enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
enumerateFromToIntegral :: forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> Stream m a
enumerateFromToIntegral a
from a
to =
    forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile (forall a. Ord a => a -> a -> Bool
<= a
to) forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *).
(Integral a, Monad m) =>
a -> a -> Stream m a
enumerateFromStepIntegral a
from a
1

{-# INLINE enumerateFromIntegral #-}
enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
enumerateFromIntegral :: forall (m :: * -> *) a.
(Monad m, Integral a, Bounded a) =>
a -> Stream m a
enumerateFromIntegral a
from = forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> Stream m a
enumerateFromToIntegral a
from forall a. Bounded a => a
maxBound

data EnumState a = EnumInit | EnumYield a a a | EnumStop

{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
enumerateFromThenToIntegralUp
    :: (Monad m, Integral a)
    => a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp :: forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp a
from a
next a
to = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p}.
Monad m =>
p -> EnumState a -> m (Step (EnumState a) a)
step forall a. EnumState a
EnumInit
    where
    {-# INLINE_LATE step #-}
    step :: p -> EnumState a -> m (Step (EnumState a) a)
step p
_ EnumState a
EnumInit =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            if a
to forall a. Ord a => a -> a -> Bool
< a
next
            then if a
to forall a. Ord a => a -> a -> Bool
< a
from
                 then forall s a. Step s a
Stop
                 else forall s a. a -> s -> Step s a
Yield a
from forall a. EnumState a
EnumStop
            else -- from <= next <= to
                let stride :: a
stride = a
next forall a. Num a => a -> a -> a
- a
from
                in forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> a -> a -> EnumState a
EnumYield a
from a
stride (a
to forall a. Num a => a -> a -> a
- a
stride)

    step p
_ (EnumYield a
x a
stride a
toMinus) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            if a
x forall a. Ord a => a -> a -> Bool
> a
toMinus
            then forall s a. a -> s -> Step s a
Yield a
x forall a. EnumState a
EnumStop
            else forall s a. a -> s -> Step s a
Yield a
x forall a b. (a -> b) -> a -> b
$ forall a. a -> a -> a -> EnumState a
EnumYield (a
x forall a. Num a => a -> a -> a
+ a
stride) a
stride a
toMinus

    step p
_ EnumState a
EnumStop = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
enumerateFromThenToIntegralDn
    :: (Monad m, Integral a)
    => a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn :: forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn a
from a
next a
to = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p}.
Monad m =>
p -> EnumState a -> m (Step (EnumState a) a)
step forall a. EnumState a
EnumInit
    where
    {-# INLINE_LATE step #-}
    step :: p -> EnumState a -> m (Step (EnumState a) a)
step p
_ EnumState a
EnumInit =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if a
to forall a. Ord a => a -> a -> Bool
> a
next
            then if a
to forall a. Ord a => a -> a -> Bool
> a
from
                 then forall s a. Step s a
Stop
                 else forall s a. a -> s -> Step s a
Yield a
from forall a. EnumState a
EnumStop
            else -- from >= next >= to
                let stride :: a
stride = a
next forall a. Num a => a -> a -> a
- a
from
                in forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> a -> a -> EnumState a
EnumYield a
from a
stride (a
to forall a. Num a => a -> a -> a
- a
stride)

    step p
_ (EnumYield a
x a
stride a
toMinus) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            if a
x forall a. Ord a => a -> a -> Bool
< a
toMinus
            then forall s a. a -> s -> Step s a
Yield a
x forall a. EnumState a
EnumStop
            else forall s a. a -> s -> Step s a
Yield a
x forall a b. (a -> b) -> a -> b
$ forall a. a -> a -> a -> EnumState a
EnumYield (a
x forall a. Num a => a -> a -> a
+ a
stride) a
stride a
toMinus

    step p
_ EnumState a
EnumStop = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
enumerateFromThenToIntegral
    :: (Monad m, Integral a)
    => a -> a -> a -> Stream m a
enumerateFromThenToIntegral :: forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegral a
from a
next a
to
    | a
next forall a. Ord a => a -> a -> Bool
>= a
from = forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp a
from a
next a
to
    | Bool
otherwise    = forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn a
from a
next a
to

{-# INLINE_NORMAL enumerateFromThenIntegral #-}
enumerateFromThenIntegral
    :: (Monad m, Integral a, Bounded a)
    => a -> a -> Stream m a
enumerateFromThenIntegral :: forall (m :: * -> *) a.
(Monad m, Integral a, Bounded a) =>
a -> a -> Stream m a
enumerateFromThenIntegral a
from a
next =
    if a
next forall a. Ord a => a -> a -> Bool
> a
from
    then forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralUp a
from a
next forall a. Bounded a => a
maxBound
    else forall (m :: * -> *) a.
(Monad m, Integral a) =>
a -> a -> a -> Stream m a
enumerateFromThenToIntegralDn a
from a
next forall a. Bounded a => a
minBound

-- For floating point numbers if the increment is less than the precision then
-- it just gets lost. Therefore we cannot always increment it correctly by just
-- repeated addition.
-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
-- 9007199254740992 + 2     :: Double => 9.007199254740994e15

-- Instead we accumulate the increment counter and compute the increment
-- every time before adding it to the starting number.
--
-- This works for Integrals as well as floating point numbers, but
-- enumerateFromStepIntegral is faster for integrals.
{-# INLINE_NORMAL enumerateFromStepNum #-}
enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum :: forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum a
from a
stride = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p}. Monad m => p -> a -> m (Step a a)
step a
0
    where
    {-# INLINE_LATE step #-}
    step :: p -> a -> m (Step a a)
step p
_ !a
i = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ (forall s a. a -> s -> Step s a
Yield forall a b. (a -> b) -> a -> b
$! (a
from forall a. Num a => a -> a -> a
+ a
i forall a. Num a => a -> a -> a
* a
stride)) forall a b. (a -> b) -> a -> b
$! (a
i forall a. Num a => a -> a -> a
+ a
1)

{-# INLINE_NORMAL numFrom #-}
numFrom :: (Monad m, Num a) => a -> Stream m a
numFrom :: forall (m :: * -> *) a. (Monad m, Num a) => a -> Stream m a
numFrom a
from = forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum a
from a
1

{-# INLINE_NORMAL numFromThen #-}
numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
numFromThen :: forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
numFromThen a
from a
next = forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum a
from (a
next forall a. Num a => a -> a -> a
- a
from)

-- We cannot write a general function for Num.  The only way to write code
-- portable between the two is to use a 'Real' constraint and convert between
-- Fractional and Integral using fromRational which is horribly slow.
{-# INLINE_NORMAL enumerateFromToFractional #-}
enumerateFromToFractional
    :: (Monad m, Fractional a, Ord a)
    => a -> a -> Stream m a
enumerateFromToFractional :: forall (m :: * -> *) a.
(Monad m, Fractional a, Ord a) =>
a -> a -> Stream m a
enumerateFromToFractional a
from a
to =
    forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile (forall a. Ord a => a -> a -> Bool
<= a
to forall a. Num a => a -> a -> a
+ a
1 forall a. Fractional a => a -> a -> a
/ a
2) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
enumerateFromStepNum a
from a
1

{-# INLINE_NORMAL enumerateFromThenToFractional #-}
enumerateFromThenToFractional
    :: (Monad m, Fractional a, Ord a)
    => a -> a -> a -> Stream m a
enumerateFromThenToFractional :: forall (m :: * -> *) a.
(Monad m, Fractional a, Ord a) =>
a -> a -> a -> Stream m a
enumerateFromThenToFractional a
from a
next a
to =
    forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile a -> Bool
predicate forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (Monad m, Num a) => a -> a -> Stream m a
numFromThen a
from a
next
    where
    mid :: a
mid = (a
next forall a. Num a => a -> a -> a
- a
from) forall a. Fractional a => a -> a -> a
/ a
2
    predicate :: a -> Bool
predicate | a
next forall a. Ord a => a -> a -> Bool
>= a
from  = (forall a. Ord a => a -> a -> Bool
<= a
to forall a. Num a => a -> a -> a
+ a
mid)
              | Bool
otherwise     = (forall a. Ord a => a -> a -> Bool
>= a
to forall a. Num a => a -> a -> a
+ a
mid)

-------------------------------------------------------------------------------
-- Generation by Conversion
-------------------------------------------------------------------------------

{-# INLINE_NORMAL fromIndicesM #-}
fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
fromIndicesM :: forall (m :: * -> *) a. Monad m => (Int -> m a) -> Stream m a
fromIndicesM Int -> m a
gen = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> Int -> m (Step Int a)
step Int
0
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i = do
       a
x <- Int -> m a
gen Int
i
       forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (Int
i forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE fromIndices #-}
fromIndices :: Monad m => (Int -> a) -> Stream m a
fromIndices :: forall (m :: * -> *) a. Monad m => (Int -> a) -> Stream m a
fromIndices Int -> a
gen = forall (m :: * -> *) a. Monad m => (Int -> m a) -> Stream m a
fromIndicesM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
gen)

{-# INLINE_NORMAL generateM #-}
generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
generateM :: forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> Stream m a
generateM Int
n Int -> m a
gen = Int
n seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> Int -> m (Step Int a)
step Int
0
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i | Int
i forall a. Ord a => a -> a -> Bool
< Int
n     = do
                           a
x <- Int -> m a
gen Int
i
                           forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (Int
i forall a. Num a => a -> a -> a
+ Int
1)
             | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE generate #-}
generate :: Monad m => Int -> (Int -> a) -> Stream m a
generate :: forall (m :: * -> *) a. Monad m => Int -> (Int -> a) -> Stream m a
generate Int
n Int -> a
gen = forall (m :: * -> *) a.
Monad m =>
Int -> (Int -> m a) -> Stream m a
generateM Int
n (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> a
gen)

-- XXX we need the MonadAsync constraint because of a rewrite rule.
-- | Convert a list of monadic actions to a 'Stream'
{-# INLINE_LATE fromListM #-}
fromListM :: MonadAsync m => [m a] -> Stream m a
fromListM :: forall (m :: * -> *) a. MonadAsync m => [m a] -> Stream m a
fromListM = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p} {a}.
Monad m =>
p -> [m a] -> m (Step [m a] a)
step
  where
    {-# INLINE_LATE step #-}
    step :: p -> [m a] -> m (Step [m a] a)
step p
_ (m a
m:[m a]
ms) = m a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x [m a]
ms
    step p
_ []     = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE toStreamD #-}
toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
toStreamD :: forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
t m a -> Stream m a
toStreamD = forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
fromStreamK forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
IsStream t =>
t m a -> Stream m a
K.toStream

{-# INLINE_NORMAL fromPrimVar #-}
fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
fromPrimVar :: forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Var IO a -> Stream m a
fromPrimVar Var IO a
var = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p}. MonadIO m => p -> () -> m (Step () a)
step ()
  where
    {-# INLINE_LATE step #-}
    step :: p -> () -> m (Step () a)
step p
_ () = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall (m :: * -> *) a. (MonadMut m, Prim a) => Var m a -> m a
readVar Var IO a
var) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
x -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x ()

-------------------------------------------------------------------------------
-- Generation from SVar
-------------------------------------------------------------------------------

data FromSVarState t m a =
      FromSVarInit
    | FromSVarRead (SVar t m a)
    | FromSVarLoop (SVar t m a) [ChildEvent a]
    | FromSVarDone (SVar t m a)

{-# INLINE_NORMAL fromSVar #-}
fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
fromSVar :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVar t m a -> Stream m a
fromSVar SVar t m a
svar = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}.
p -> FromSVarState t m a -> m (Step (FromSVarState t m a) a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
FromSVarState t m a
FromSVarInit
    where

    {-# INLINE_LATE step #-}
    step :: p -> FromSVarState t m a -> m (Step (FromSVarState t m a) a)
step p
_ FromSVarState t m a
FromSVarInit = do
        IORef ()
ref <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef ()
        Weak (IORef ())
_ <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ()
ref IO ()
hook
        -- when this copy of svar gets garbage collected "ref" will get
        -- garbage collected and our GC hook will be called.
        let sv :: SVar t m a
sv = SVar t m a
svar{svarRef :: Maybe (IORef ())
svarRef = forall a. a -> Maybe a
Just IORef ()
ref}
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> FromSVarState t m a
FromSVarRead SVar t m a
sv)

        where

        {-# NOINLINE hook #-}
        hook :: IO ()
hook = do
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> Bool
svarInspectMode SVar t m a
svar) forall a b. (a -> b) -> a -> b
$ do
                Maybe AbsTime
r <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef (SVarStats -> IORef (Maybe AbsTime)
svarStopTime (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SVarStats
svarStats SVar t m a
svar))
                forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Maybe a -> Bool
isNothing Maybe AbsTime
r) forall a b. (a -> b) -> a -> b
$
                    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> String -> IO ()
printSVar SVar t m a
svar String
"SVar Garbage Collected"
            forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> IO ()
cleanupSVar SVar t m a
svar
            -- If there are any SVars referenced by this SVar a GC will prompt
            -- them to be cleaned up quickly.
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> Bool
svarInspectMode SVar t m a
svar) IO ()
performMajorGC

    step p
_ (FromSVarRead SVar t m a
sv) = do
        [ChildEvent a]
list <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> m [ChildEvent a]
readOutputQ SVar t m a
sv
        -- Reversing the output is important to guarantee that we process the
        -- outputs in the same order as they were generated by the constituent
        -- streams.
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> [ChildEvent a] -> FromSVarState t m a
FromSVarLoop SVar t m a
sv (forall a. [a] -> [a]
Prelude.reverse [ChildEvent a]
list)

    step p
_ (FromSVarLoop SVar t m a
sv []) = do
        Bool
done <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> m Bool
postProcess SVar t m a
sv
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ if Bool
done
                      then (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> FromSVarState t m a
FromSVarDone SVar t m a
sv)
                      else (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> FromSVarState t m a
FromSVarRead SVar t m a
sv)

    step p
_ (FromSVarLoop SVar t m a
sv (ChildEvent a
ev : [ChildEvent a]
es)) = do
        case ChildEvent a
ev of
            ChildYield a
a -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> [ChildEvent a] -> FromSVarState t m a
FromSVarLoop SVar t m a
sv [ChildEvent a]
es)
            ChildStop ThreadId
tid Maybe SomeException
e -> do
                forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> ThreadId -> m ()
accountThread SVar t m a
sv ThreadId
tid
                case Maybe SomeException
e of
                    Maybe SomeException
Nothing -> do
                        Bool
stop <- forall {m :: * -> *}. MonadIO m => ThreadId -> m Bool
shouldStop ThreadId
tid
                        if Bool
stop
                        then do
                            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> IO ()
cleanupSVar SVar t m a
sv)
                            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> FromSVarState t m a
FromSVarDone SVar t m a
sv)
                        else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> [ChildEvent a] -> FromSVarState t m a
FromSVarLoop SVar t m a
sv [ChildEvent a]
es)
                    Just SomeException
ex ->
                        case forall e. Exception e => SomeException -> Maybe e
fromException SomeException
ex of
                            Just ThreadAbort
ThreadAbort ->
                                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> [ChildEvent a] -> FromSVarState t m a
FromSVarLoop SVar t m a
sv [ChildEvent a]
es)
                            Maybe ThreadAbort
Nothing -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> IO ()
cleanupSVar SVar t m a
sv) forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM SomeException
ex
        where

        shouldStop :: ThreadId -> m Bool
shouldStop ThreadId
tid =
            case forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SVarStopStyle
svarStopStyle SVar t m a
sv of
                SVarStopStyle
StopNone -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
                SVarStopStyle
StopAny -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
                SVarStopStyle
StopBy -> do
                    ThreadId
sid <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO a
readIORef (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> IORef ThreadId
svarStopBy SVar t m a
sv)
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if ThreadId
tid forall a. Eq a => a -> a -> Bool
== ThreadId
sid then Bool
True else Bool
False

    step p
_ (FromSVarDone SVar t m a
sv) = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> Bool
svarInspectMode SVar t m a
sv) forall a b. (a -> b) -> a -> b
$ do
            AbsTime
t <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Clock -> IO AbsTime
getTime Clock
Monotonic
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef (SVarStats -> IORef (Maybe AbsTime)
svarStopTime (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SVarStats
svarStats SVar t m a
sv)) (forall a. a -> Maybe a
Just AbsTime
t)
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> String -> IO ()
printSVar SVar t m a
sv String
"SVar Done"
        forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-------------------------------------------------------------------------------
-- Process events received by a fold consumer from a stream producer
-------------------------------------------------------------------------------

{-# INLINE_NORMAL fromProducer #-}
fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
fromProducer :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVar t m a -> Stream m a
fromProducer SVar t m a
svar = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p} {t :: (* -> *) -> * -> *} {a}.
MonadIO m =>
p -> FromSVarState t m a -> m (Step (FromSVarState t m a) a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> FromSVarState t m a
FromSVarRead SVar t m a
svar)
    where

    {-# INLINE_LATE step #-}
    step :: p -> FromSVarState t m a -> m (Step (FromSVarState t m a) a)
step p
_ (FromSVarRead SVar t m a
sv) = do
        [ChildEvent a]
list <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> m [ChildEvent a]
readOutputQ SVar t m a
sv
        -- Reversing the output is important to guarantee that we process the
        -- outputs in the same order as they were generated by the constituent
        -- streams.
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> [ChildEvent a] -> FromSVarState t m a
FromSVarLoop SVar t m a
sv (forall a. [a] -> [a]
Prelude.reverse [ChildEvent a]
list)

    step p
_ (FromSVarLoop SVar t m a
sv []) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> FromSVarState t m a
FromSVarRead SVar t m a
sv
    step p
_ (FromSVarLoop SVar t m a
sv (ChildEvent a
ev : [ChildEvent a]
es)) = do
        case ChildEvent a
ev of
            ChildYield a
a -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> [ChildEvent a] -> FromSVarState t m a
FromSVarLoop SVar t m a
sv [ChildEvent a]
es)
            ChildStop ThreadId
tid Maybe SomeException
e -> do
                forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> ThreadId -> m ()
accountThread SVar t m a
sv ThreadId
tid
                case Maybe SomeException
e of
                    Maybe SomeException
Nothing -> do
                        forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> m ()
sendStopToProducer SVar t m a
sv
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> FromSVarState t m a
FromSVarDone SVar t m a
sv)
                    Just SomeException
_ -> forall a. HasCallStack => String -> a
error String
"Bug: fromProducer: received exception"

    step p
_ (FromSVarDone SVar t m a
sv) = do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> Bool
svarInspectMode SVar t m a
sv) forall a b. (a -> b) -> a -> b
$ do
            AbsTime
t <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Clock -> IO AbsTime
getTime Clock
Monotonic
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef (SVarStats -> IORef (Maybe AbsTime)
svarStopTime (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SVarStats
svarStats SVar t m a
sv)) (forall a. a -> Maybe a
Just AbsTime
t)
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> String -> IO ()
printSVar SVar t m a
sv String
"SVar Done"
        forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step p
_ FromSVarState t m a
FromSVarInit = forall a. HasCallStack => a
undefined

-------------------------------------------------------------------------------
-- Hoisting the inner monad
-------------------------------------------------------------------------------

{-# INLINE_NORMAL hoist #-}
hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
hoist :: forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> Stream m a -> Stream n a
hoist forall x. m x -> n x
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}. State Stream m a -> s -> n (Step s a)
step' s
state)
    where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> n (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
r <- forall x. m x -> n x
f forall a b. (a -> b) -> a -> b
$ State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x s
s
            Skip  s
s   -> forall s a. s -> Step s a
Skip s
s
            Step s a
Stop      -> forall s a. Step s a
Stop

{-# INLINE generally #-}
generally :: Monad m => Stream Identity a -> Stream m a
generally :: forall (m :: * -> *) a. Monad m => Stream Identity a -> Stream m a
generally = forall (n :: * -> *) (m :: * -> *) a.
Monad n =>
(forall x. m x -> n x) -> Stream m a -> Stream n a
hoist (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Identity a -> a
runIdentity)

{-# INLINE_NORMAL liftInner #-}
liftInner :: (Monad m, MonadTrans t, Monad (t m))
    => Stream m a -> Stream (t m) a
liftInner :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
(Monad m, MonadTrans t, Monad (t m)) =>
Stream m a -> Stream (t m) a
liftInner (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {t :: (* -> *) -> * -> *} {m :: * -> *} {a}.
(MonadTrans t, Monad (t m)) =>
State Stream m a -> s -> t m (Step s a)
step' s
state
    where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> t m (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x s
s
            Skip s
s    -> forall s a. s -> Step s a
Skip s
s
            Step s a
Stop      -> forall s a. Step s a
Stop

{-# INLINE_NORMAL runReaderT #-}
runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
runReaderT :: forall (m :: * -> *) s a.
Monad m =>
s -> Stream (ReaderT s m) a -> Stream m a
runReaderT s
sval (Stream State Stream (ReaderT s m) a -> s -> ReaderT s m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}. State Stream m a -> s -> m (Step s a)
step' s
state
    where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
r <- forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
Reader.runReaderT (State Stream (ReaderT s m) a -> s -> ReaderT s m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st) s
sval
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x s
s
            Skip  s
s   -> forall s a. s -> Step s a
Skip s
s
            Step s a
Stop      -> forall s a. Step s a
Stop

{-# INLINE_NORMAL evalStateT #-}
evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
evalStateT :: forall (m :: * -> *) s a.
Monad m =>
s -> Stream (StateT s m) a -> Stream m a
evalStateT s
sval (Stream State Stream (StateT s m) a -> s -> StateT s m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a -> (s, s) -> m (Step (s, s) a)
step' (s
state, s
sval)
    where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, s) -> m (Step (s, s) a)
step' State Stream m a
gst (s
st, s
sv) = do
        (Step s a
r, s
sv') <- forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
State.runStateT (State Stream (StateT s m) a -> s -> StateT s m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st) s
sv
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x (s
s, s
sv')
            Skip  s
s   -> forall s a. s -> Step s a
Skip (s
s, s
sv')
            Step s a
Stop      -> forall s a. Step s a
Stop

{-# INLINE_NORMAL runStateT #-}
runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
runStateT :: forall (m :: * -> *) s a.
Monad m =>
s -> Stream (StateT s m) a -> Stream m (s, a)
runStateT s
sval (Stream State Stream (StateT s m) a -> s -> StateT s m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a -> (s, s) -> m (Step (s, s) (s, a))
step' (s
state, s
sval)
    where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, s) -> m (Step (s, s) (s, a))
step' State Stream m a
gst (s
st, s
sv) = do
        (Step s a
r, s
sv') <- forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
State.runStateT (State Stream (StateT s m) a -> s -> StateT s m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st) s
sv
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield (s
sv', a
x) (s
s, s
sv')
            Skip  s
s   -> forall s a. s -> Step s a
Skip (s
s, s
sv')
            Step s a
Stop      -> forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Elimination by Folds
------------------------------------------------------------------------------

------------------------------------------------------------------------------
-- Right Folds
------------------------------------------------------------------------------

{-# INLINE_NORMAL foldr1 #-}
foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
foldr1 :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> a) -> Stream m a -> m (Maybe a)
foldr1 a -> a -> a
f Stream m a
m = do
     Maybe (a, Stream m a)
r <- forall (m :: * -> *) a.
Monad m =>
Stream m a -> m (Maybe (a, Stream m a))
uncons Stream m a
m
     case Maybe (a, Stream m a)
r of
         Maybe (a, Stream m a)
Nothing   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
         Just (a
h, Stream m a
t) -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. a -> Maybe a
Just (forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
foldr a -> a -> a
f a
h Stream m a
t)

------------------------------------------------------------------------------
-- Left Folds
------------------------------------------------------------------------------

{-# INLINE_NORMAL foldlT #-}
foldlT :: (Monad m, Monad (s m), MonadTrans s)
    => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
foldlT :: forall (m :: * -> *) (s :: (* -> *) -> * -> *) b a.
(Monad m, Monad (s m), MonadTrans s) =>
(s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
foldlT s m b -> a -> s m b
fstep s m b
begin (Stream State Stream m a -> s -> m (Step s a)
step s
state) = SPEC -> s m b -> s -> s m b
go SPEC
SPEC s m b
begin s
state
  where
    go :: SPEC -> s m b -> s -> s m b
go !SPEC
_ s m b
acc s
st = do
        Step s a
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> SPEC -> s m b -> s -> s m b
go SPEC
SPEC (s m b -> a -> s m b
fstep s m b
acc a
x) s
s
            Skip s
s -> SPEC -> s m b -> s -> s m b
go SPEC
SPEC s m b
acc s
s
            Step s a
Stop   -> s m b
acc

-- Note, this is going to have horrible performance, because of the nature of
-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
-- likely be practically unusable.
{-# INLINE_NORMAL foldlS #-}
foldlS :: Monad m
    => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
foldlS :: forall (m :: * -> *) b a.
Monad m =>
(Stream m b -> a -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldlS Stream m b -> a -> Stream m b
fstep Stream m b
begin (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> Either (s, Stream m b) (Stream m b)
-> m (Step (Either (s, Stream m b) (Stream m b)) b)
step' (forall a b. a -> Either a b
Left (s
state, Stream m b
begin))
  where
    step' :: State Stream m a
-> Either (s, Stream m b) (Stream m b)
-> m (Step (Either (s, Stream m b) (Stream m b)) b)
step' State Stream m a
gst (Left (s
st, Stream m b
acc)) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left (s
s, Stream m b -> a -> Stream m b
fstep Stream m b
acc a
x))
            Skip s
s -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left (s
s, Stream m b
acc))
            Step s a
Stop   -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right Stream m b
acc)

    step' State Stream m a
gst (Right (Stream State Stream m b -> s -> m (Step s b)
stp s
stt)) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
stp (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall a b. b -> Either a b
Right (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
stp s
s))
            Skip s
s -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
stp s
s))
            Step s b
Stop   -> forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Parses
------------------------------------------------------------------------------

-- Inlined definition. Without the inline "serially/parser/take" benchmark
-- degrades and splitParse does not fuse. Even using "inline" at the callsite
-- does not help.
{-# INLINE splitAt #-}
splitAt :: Int -> [a] -> ([a],[a])
splitAt :: forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
ls
  | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = ([], [a]
ls)
  | Bool
otherwise          = forall a. Int -> [a] -> ([a], [a])
splitAt' Int
n [a]
ls
    where
        splitAt' :: Int -> [a] -> ([a], [a])
        splitAt' :: forall a. Int -> [a] -> ([a], [a])
splitAt' Int
_  []     = ([], [])
        splitAt' Int
1  (a
x:[a]
xs) = ([a
x], [a]
xs)
        splitAt' Int
m  (a
x:[a]
xs) = (a
xforall a. a -> [a] -> [a]
:[a]
xs', [a]
xs'')
          where
            ([a]
xs', [a]
xs'') = forall a. Int -> [a] -> ([a], [a])
splitAt' (Int
m forall a. Num a => a -> a -> a
- Int
1) [a]
xs

-- | Run a 'Parse' over a stream.
{-# INLINE_NORMAL parselMx' #-}
parselMx'
    :: MonadThrow m
    => (s -> a -> m (PR.Step s b))
    -> m s
    -> (s -> m b)
    -> Stream m a
    -> m b
parselMx' :: forall (m :: * -> *) s a b.
MonadThrow m =>
(s -> a -> m (Step s b)) -> m s -> (s -> m b) -> Stream m a -> m b
parselMx' s -> a -> m (Step s b)
pstep m s
initial s -> m b
extract (Stream State Stream m a -> s -> m (Step s a)
step s
state) = do
    m s
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> s -> [a] -> s -> m b
go SPEC
SPEC s
state []

    where

    -- XXX currently we are using a dumb list based approach for backtracking
    -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
    -- That will allow us more efficient random back and forth movement.
    {-# INLINE go #-}
    go :: SPEC -> s -> [a] -> s -> m b
go !SPEC
_ s
st [a]
buf !s
pst = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Step s b
pRes <- s -> a -> m (Step s b)
pstep s
pst a
x
                case Step s b
pRes of
                    -- PR.Yield 0 pst1 -> go SPEC s [] pst1
                    PR.Yield Int
n s
pst1 -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        SPEC -> s -> [a] -> s -> m b
go SPEC
SPEC s
s (forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)) s
pst1
                    PR.Skip Int
0 s
pst1 -> SPEC -> s -> [a] -> s -> m b
go SPEC
SPEC s
s (a
xforall a. a -> [a] -> [a]
:[a]
buf) s
pst1
                    PR.Skip Int
n s
pst1 -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        let ([a]
src0, [a]
buf1) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)
                            src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0
                        SPEC -> s -> [a] -> [a] -> s -> m b
gobuf SPEC
SPEC s
s [a]
buf1 [a]
src s
pst1
                    PR.Stop Int
_ b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return b
b
                    PR.Error String
err -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ String -> ParseError
ParseError String
err
            Skip s
s -> SPEC -> s -> [a] -> s -> m b
go SPEC
SPEC s
s [a]
buf s
pst
            Step s a
Stop   -> s -> m b
extract s
pst

    gobuf :: SPEC -> s -> [a] -> [a] -> s -> m b
gobuf !SPEC
_ s
s [a]
buf [] !s
pst = SPEC -> s -> [a] -> s -> m b
go SPEC
SPEC s
s [a]
buf s
pst
    gobuf !SPEC
_ s
s [a]
buf (a
x:[a]
xs) !s
pst = do
        Step s b
pRes <- s -> a -> m (Step s b)
pstep s
pst a
x
        case Step s b
pRes of
            -- PR.Yield 0 pst1 -> go SPEC s [] pst1
            PR.Yield Int
n s
pst1 -> do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                SPEC -> s -> [a] -> [a] -> s -> m b
gobuf SPEC
SPEC s
s (forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)) [a]
xs s
pst1
            PR.Skip Int
0 s
pst1 -> SPEC -> s -> [a] -> [a] -> s -> m b
gobuf SPEC
SPEC s
s (a
xforall a. a -> [a] -> [a]
:[a]
buf) [a]
xs s
pst1
            PR.Skip Int
n s
pst1 -> do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let ([a]
src0, [a]
buf1) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)
                    src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0 forall a. [a] -> [a] -> [a]
++ [a]
xs
                SPEC -> s -> [a] -> [a] -> s -> m b
gobuf SPEC
SPEC s
s [a]
buf1 [a]
src s
pst1
            PR.Stop Int
_ b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return b
b
            PR.Error String
err -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ String -> ParseError
ParseError String
err

------------------------------------------------------------------------------
-- Repeated parsing
------------------------------------------------------------------------------

{-# ANN type ParseChunksState Fuse #-}
data ParseChunksState x inpBuf st pst =
      ParseChunksInit inpBuf st
    | ParseChunksInitLeftOver inpBuf
    | ParseChunksStream st inpBuf pst
    | ParseChunksBuf inpBuf st inpBuf pst
    | ParseChunksYield x (ParseChunksState x inpBuf st pst)

{-# INLINE_NORMAL splitParse #-}
splitParse
    :: MonadThrow m
    => Parser m a b
    -> Stream m a
    -> Stream m b
splitParse :: forall (m :: * -> *) a b.
MonadThrow m =>
Parser m a b -> Stream m a -> Stream m b
splitParse (Parser s -> a -> m (Step s b)
pstep m s
initial s -> m b
extract) (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> ParseChunksState b [a] s s
-> m (Step (ParseChunksState b [a] s s) b)
stepOuter (forall x inpBuf st pst.
inpBuf -> st -> ParseChunksState x inpBuf st pst
ParseChunksInit [] s
state)

    where

    {-# INLINE_LATE stepOuter #-}
    -- Buffer is empty, go to stream processing loop
    stepOuter :: State Stream m a
-> ParseChunksState b [a] s s
-> m (Step (ParseChunksState b [a] s s) b)
stepOuter State Stream m a
_ (ParseChunksInit [] s
st) = do
        m s
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. s -> Step s a
Skip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x inpBuf st pst.
st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksStream s
st []

    -- Buffer is not empty, go to buffered processing loop
    stepOuter State Stream m a
_ (ParseChunksInit [a]
src s
st) = do
        m s
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s a. s -> Step s a
Skip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall x inpBuf st pst.
inpBuf -> st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksBuf [a]
src s
st []

    -- XXX we just discard any leftover input at the end
    stepOuter State Stream m a
_ (ParseChunksInitLeftOver [a]
_) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    -- Buffer is empty process elements from the stream
    stepOuter State Stream m a
gst (ParseChunksStream s
st [a]
buf s
pst) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Step s b
pRes <- s -> a -> m (Step s b)
pstep s
pst a
x
                case Step s b
pRes of
                    -- PR.Yield 0 pst1 -> go SPEC s [] pst1
                    PR.Yield Int
n s
pst1 -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        let buf1 :: [a]
buf1 = forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x inpBuf st pst.
st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksStream s
s [a]
buf1 s
pst1
                    -- PR.Skip 0 pst1 ->
                    --     return $ Skip $ ParseChunksStream s (x:buf) pst1
                    PR.Skip Int
n s
pst1 -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        let ([a]
src0, [a]
buf1) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)
                            src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x inpBuf st pst.
inpBuf -> st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksBuf [a]
src s
s [a]
buf1 s
pst1
                    -- XXX Specialize for Stop 0 common case?
                    PR.Stop Int
n b
b -> do
                        forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                        let src :: [a]
src = forall a. [a] -> [a]
Prelude.reverse (forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf))
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$
                            forall x inpBuf st pst.
x
-> ParseChunksState x inpBuf st pst
-> ParseChunksState x inpBuf st pst
ParseChunksYield b
b (forall x inpBuf st pst.
inpBuf -> st -> ParseChunksState x inpBuf st pst
ParseChunksInit [a]
src s
s)
                    PR.Error String
err -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ String -> ParseError
ParseError String
err
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x inpBuf st pst.
st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksStream s
s [a]
buf s
pst
            Step s a
Stop   -> do
                b
b <- s -> m b
extract s
pst
                let src :: [a]
src = forall a. [a] -> [a]
Prelude.reverse [a]
buf
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$
                    forall x inpBuf st pst.
x
-> ParseChunksState x inpBuf st pst
-> ParseChunksState x inpBuf st pst
ParseChunksYield b
b (forall x inpBuf st pst. inpBuf -> ParseChunksState x inpBuf st pst
ParseChunksInitLeftOver [a]
src)

    -- go back to stream processing mode
    stepOuter State Stream m a
_ (ParseChunksBuf [] s
s [a]
buf s
pst) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x inpBuf st pst.
st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksStream s
s [a]
buf s
pst

    -- buffered processing loop
    stepOuter State Stream m a
_ (ParseChunksBuf (a
x:[a]
xs) s
s [a]
buf s
pst) = do
        Step s b
pRes <- s -> a -> m (Step s b)
pstep s
pst a
x
        case Step s b
pRes of
            -- PR.Yield 0 pst1 ->
            PR.Yield Int
n s
pst1 ->  do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let buf1 :: [a]
buf1 = forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x inpBuf st pst.
inpBuf -> st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksBuf [a]
xs s
s [a]
buf1 s
pst1
         -- PR.Skip 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
            PR.Skip Int
n s
pst1 -> do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let ([a]
src0, [a]
buf1) = forall a. Int -> [a] -> ([a], [a])
splitAt Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)
                    src :: [a]
src  = forall a. [a] -> [a]
Prelude.reverse [a]
src0 forall a. [a] -> [a] -> [a]
++ [a]
xs
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x inpBuf st pst.
inpBuf -> st -> inpBuf -> pst -> ParseChunksState x inpBuf st pst
ParseChunksBuf [a]
src s
s [a]
buf1 s
pst1
            -- XXX Specialize for Stop 0 common case?
            PR.Stop Int
n b
b -> do
                forall a. HasCallStack => Bool -> a -> a
assert (Int
n forall a. Ord a => a -> a -> Bool
<= forall (t :: * -> *) a. Foldable t => t a -> Int
length (a
xforall a. a -> [a] -> [a]
:[a]
buf)) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
                let src :: [a]
src = forall a. [a] -> [a]
Prelude.reverse (forall a. Int -> [a] -> [a]
Prelude.take Int
n (a
xforall a. a -> [a] -> [a]
:[a]
buf)) forall a. [a] -> [a] -> [a]
++ [a]
xs
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x inpBuf st pst.
x
-> ParseChunksState x inpBuf st pst
-> ParseChunksState x inpBuf st pst
ParseChunksYield b
b (forall x inpBuf st pst.
inpBuf -> st -> ParseChunksState x inpBuf st pst
ParseChunksInit [a]
src s
s)
            PR.Error String
err -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM forall a b. (a -> b) -> a -> b
$ String -> ParseError
ParseError String
err

    stepOuter State Stream m a
_ (ParseChunksYield b
a ParseChunksState b [a] s s
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
a ParseChunksState b [a] s s
next

------------------------------------------------------------------------------
-- Specialized Folds
------------------------------------------------------------------------------

-- | Run a streaming composition, discard the results.
{-# INLINE_LATE drain #-}
drain :: Monad m => Stream m a -> m ()
-- drain = foldrM (\_ xs -> xs) (return ())
drain :: forall (m :: * -> *) a. Monad m => Stream m a -> m ()
drain (Stream State Stream m a -> s -> m (Step s a)
step s
state) = SPEC -> s -> m ()
go SPEC
SPEC s
state
  where
    go :: SPEC -> s -> m ()
go !SPEC
_ s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
_ s
s -> SPEC -> s -> m ()
go SPEC
SPEC s
s
            Skip s
s    -> SPEC -> s -> m ()
go SPEC
SPEC s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return ()

{-# INLINE_NORMAL null #-}
null :: Monad m => Stream m a -> m Bool
null :: forall (m :: * -> *) a. Monad m => Stream m a -> m Bool
null Stream m a
m = forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
_ m Bool
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False) (forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True) Stream m a
m

{-# INLINE_NORMAL head #-}
head :: Monad m => Stream m a -> m (Maybe a)
head :: forall (m :: * -> *) a. Monad m => Stream m a -> m (Maybe a)
head Stream m a
m = forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
x m (Maybe a)
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
x)) (forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing) Stream m a
m

{-# INLINE_NORMAL headElse #-}
headElse :: Monad m => a -> Stream m a -> m a
headElse :: forall (m :: * -> *) a. Monad m => a -> Stream m a -> m a
headElse a
a Stream m a
m = forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
x m a
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return a
x) (forall (m :: * -> *) a. Monad m => a -> m a
return a
a) Stream m a
m

-- Does not fuse, has the same performance as the StreamK version.
{-# INLINE_NORMAL tail #-}
tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
tail :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> m (Maybe (Stream m a))
tail (UnStream State Stream m a -> s -> m (Step s a)
step s
state) = s -> m (Maybe (Stream m a))
go s
state
  where
    go :: s -> m (Maybe (Stream m a))
go s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
_ s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step s
s)
            Skip  s
s   -> s -> m (Maybe (Stream m a))
go s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

-- XXX will it fuse? need custom impl?
{-# INLINE_NORMAL last #-}
last :: Monad m => Stream m a -> m (Maybe a)
last :: forall (m :: * -> *) a. Monad m => Stream m a -> m (Maybe a)
last = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
foldl' (\Maybe a
_ a
y -> forall a. a -> Maybe a
Just a
y) forall a. Maybe a
Nothing

{-# INLINE_NORMAL elem #-}
elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
elem :: forall (m :: * -> *) a.
(Monad m, Eq a) =>
a -> Stream m a -> m Bool
elem a
e (Stream State Stream m a -> s -> m (Step s a)
step s
state) = s -> m Bool
go s
state
  where
    go :: s -> m Bool
go s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s
              | a
x forall a. Eq a => a -> a -> Bool
== a
e    -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
              | Bool
otherwise -> s -> m Bool
go s
s
            Skip s
s -> s -> m Bool
go s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

{-# INLINE_NORMAL notElem #-}
notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
notElem :: forall (m :: * -> *) a.
(Monad m, Eq a) =>
a -> Stream m a -> m Bool
notElem a
e Stream m a
s = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Bool -> Bool
not (forall (m :: * -> *) a.
(Monad m, Eq a) =>
a -> Stream m a -> m Bool
elem a
e Stream m a
s)

{-# INLINE_NORMAL all #-}
all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
all :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> m Bool
all a -> Bool
p (Stream State Stream m a -> s -> m (Step s a)
step s
state) = s -> m Bool
go s
state
  where
    go :: s -> m Bool
go s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s
              | a -> Bool
p a
x       -> s -> m Bool
go s
s
              | Bool
otherwise -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
            Skip s
s -> s -> m Bool
go s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

{-# INLINE_NORMAL any #-}
any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
any :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> m Bool
any a -> Bool
p (Stream State Stream m a -> s -> m (Step s a)
step s
state) = s -> m Bool
go s
state
  where
    go :: s -> m Bool
go s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s
              | a -> Bool
p a
x       -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
              | Bool
otherwise -> s -> m Bool
go s
s
            Skip s
s -> s -> m Bool
go s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

{-# INLINE_NORMAL maximum #-}
maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
maximum :: forall (m :: * -> *) a.
(Monad m, Ord a) =>
Stream m a -> m (Maybe a)
maximum (Stream State Stream m a -> s -> m (Step s a)
step s
state) = Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
state
  where
    go :: Maybe a -> s -> m (Maybe a)
go Maybe a
Nothing s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
            Skip  s
s   -> Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    go (Just a
acc) s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s
              | a
acc forall a. Ord a => a -> a -> Bool
<= a
x  -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
              | Bool
otherwise -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
            Skip s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
acc)

{-# INLINE_NORMAL maximumBy #-}
maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
maximumBy :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> Ordering) -> Stream m a -> m (Maybe a)
maximumBy a -> a -> Ordering
cmp (Stream State Stream m a -> s -> m (Step s a)
step s
state) = Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
state
  where
    go :: Maybe a -> s -> m (Maybe a)
go Maybe a
Nothing s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
            Skip  s
s   -> Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    go (Just a
acc) s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> case a -> a -> Ordering
cmp a
acc a
x of
                Ordering
GT -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
                Ordering
_  -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
            Skip s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
acc)

{-# INLINE_NORMAL minimum #-}
minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
minimum :: forall (m :: * -> *) a.
(Monad m, Ord a) =>
Stream m a -> m (Maybe a)
minimum (Stream State Stream m a -> s -> m (Step s a)
step s
state) = Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
state
  where
    go :: Maybe a -> s -> m (Maybe a)
go Maybe a
Nothing s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
            Skip  s
s   -> Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    go (Just a
acc) s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s
              | a
acc forall a. Ord a => a -> a -> Bool
<= a
x  -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
              | Bool
otherwise -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
            Skip s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
acc)

{-# INLINE_NORMAL minimumBy #-}
minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
minimumBy :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> Ordering) -> Stream m a -> m (Maybe a)
minimumBy a -> a -> Ordering
cmp (Stream State Stream m a -> s -> m (Step s a)
step s
state) = Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
state
  where
    go :: Maybe a -> s -> m (Maybe a)
go Maybe a
Nothing s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
            Skip  s
s   -> Maybe a -> s -> m (Maybe a)
go forall a. Maybe a
Nothing s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    go (Just a
acc) s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> case a -> a -> Ordering
cmp a
acc a
x of
                Ordering
GT -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
x) s
s
                Ordering
_  -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
            Skip s
s -> Maybe a -> s -> m (Maybe a)
go (forall a. a -> Maybe a
Just a
acc) s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
acc)

{-# INLINE_NORMAL (!!) #-}
(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
(Stream State Stream m a -> s -> m (Step s a)
step s
state) !! :: forall (m :: * -> *) a. Monad m => Stream m a -> Int -> m (Maybe a)
!! Int
i = forall {t}. (Ord t, Num t) => t -> s -> m (Maybe a)
go Int
i s
state
  where
    go :: t -> s -> m (Maybe a)
go t
n s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s | t
n forall a. Ord a => a -> a -> Bool
< t
0 -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
                      | t
n forall a. Eq a => a -> a -> Bool
== t
0 -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
x
                      | Bool
otherwise -> t -> s -> m (Maybe a)
go (t
n forall a. Num a => a -> a -> a
- t
1) s
s
            Skip s
s -> t -> s -> m (Maybe a)
go t
n s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

{-# INLINE_NORMAL lookup #-}
lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
lookup :: forall (m :: * -> *) a b.
(Monad m, Eq a) =>
a -> Stream m (a, b) -> m (Maybe b)
lookup a
e Stream m (a, b)
m = forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\(a
a, b
b) m (Maybe b)
xs -> if a
e forall a. Eq a => a -> a -> Bool
== a
a then forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just b
b) else m (Maybe b)
xs)
                   (forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing) Stream m (a, b)
m

{-# INLINE_NORMAL findM #-}
findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
findM :: forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> m (Maybe a)
findM a -> m Bool
p Stream m a
m = forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
x m (Maybe a)
xs -> a -> m Bool
p a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Bool
r -> if Bool
r then forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
x) else m (Maybe a)
xs)
                   (forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing) Stream m a
m

{-# INLINE find #-}
find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
find :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> m (Maybe a)
find a -> Bool
p = forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> m (Maybe a)
findM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
p)

{-# INLINE_NORMAL findIndices #-}
findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
findIndices :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m Int
findIndices a -> Bool
p (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {b} {m :: * -> *} {a}.
Num b =>
State Stream m a -> (s, b) -> m (Step (s, b) b)
step' (s
state, Int
0)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, b) -> m (Step (s, b) b)
step' State Stream m a
gst (s
st, b
i) = b
i seq :: forall a b. a -> b -> b
`seq` do
      Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
          Yield a
x s
s -> if a -> Bool
p a
x then forall s a. a -> s -> Step s a
Yield b
i (s
s, b
iforall a. Num a => a -> a -> a
+b
1) else forall s a. s -> Step s a
Skip (s
s, b
iforall a. Num a => a -> a -> a
+b
1)
          Skip s
s -> forall s a. s -> Step s a
Skip (s
s, b
i)
          Step s a
Stop   -> forall s a. Step s a
Stop

{-# INLINE toListRev #-}
toListRev :: Monad m => Stream m a -> m [a]
toListRev :: forall (m :: * -> *) a. Monad m => Stream m a -> m [a]
toListRev = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip (:)) []

-- We can implement reverse as:
--
-- > reverse = foldlS (flip cons) nil
--
-- However, this implementation is unusable because of the horrible performance
-- of cons. So we just convert it to a list first and then stream from the
-- list.
--
-- XXX Maybe we can use an Array instead of a list here?
{-# INLINE_NORMAL reverse #-}
reverse :: Monad m => Stream m a -> Stream m a
reverse :: forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
reverse Stream m a
m = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> Maybe [a] -> m (Step (Maybe [a]) a)
step forall a. Maybe a
Nothing
    where
    {-# INLINE_LATE step #-}
    step :: p -> Maybe [a] -> m (Step (Maybe [a]) a)
step p
_ Maybe [a]
Nothing = do
        [a]
xs <- forall (m :: * -> *) a. Monad m => Stream m a -> m [a]
toListRev Stream m a
m
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just [a]
xs)
    step p
_ (Just (a
x:[a]
xs)) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just [a]
xs)
    step p
_ (Just []) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- Much faster reverse for Storables
{-# INLINE_NORMAL reverse' #-}
reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
{-
-- This commented implementation copies the whole stream into one single array
-- and then streams from that array, this is 3-4x faster than the chunked code
-- that follows.  Though this could be problematic due to unbounded large
-- allocations. We need to figure out why the chunked code is slower and if we
-- can optimize the chunked code to work as fast as this one. It may be a
-- fusion issue?
import Foreign.ForeignPtr (touchForeignPtr)
import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
import Foreign.Ptr (Ptr, plusPtr)
reverse' m = Stream step Nothing
    where
    {-# INLINE_LATE step #-}
    step _ Nothing = do
        arr <- A.fromStreamD m
        let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
        return $ Skip $ Just (aStart arr, p)

    step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop

    step _ (Just (start, p)) = do
        let !x = A.unsafeInlineIO $ do
                    r <- peek p
                    touchForeignPtr start
                    return r
            next = p `plusPtr` negate (sizeOf (undefined :: a))
        return $ Yield x (Just (start, next))
-}
reverse' :: forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m a -> Stream m a
reverse' Stream m a
m =
          forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Stream m (Array a) -> Stream m a
A.flattenArraysRev
        forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
fromStreamK
        forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
IsStream t =>
t m a -> t m a
K.reverse
        forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
toStreamK
        forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Storable a) =>
Int -> Stream m a -> Stream m (Array a)
A.fromStreamDArraysOf Int
A.defaultChunkSize Stream m a
m


------------------------------------------------------------------------------
-- Grouping/Splitting
------------------------------------------------------------------------------

{-# INLINE_NORMAL splitSuffixBy' #-}
splitSuffixBy' :: Monad m
    => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
splitSuffixBy' :: forall (m :: * -> *) a b.
Monad m =>
(a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
splitSuffixBy' a -> Bool
predicate Fold m a b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (forall {a} {m :: * -> *} {a}.
Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
stepOuter Fold m a b
f) (forall a. a -> Maybe a
Just s
state)

    where

    {-# INLINE_LATE stepOuter #-}
    stepOuter :: Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
stepOuter (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                s
acc' <- s -> a -> m s
fstep s
acc a
x
                if (a -> Bool
predicate a
x)
                then s -> m a
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
val -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
val (forall a. a -> Maybe a
Just s
s)
                else SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc'

            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where

        go :: SPEC -> s -> s -> m (Step (Maybe s) a)
go !SPEC
_ s
stt !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
            case Step s a
res of
                Yield a
x s
s -> do
                    s
acc' <- s -> a -> m s
fstep s
acc a
x
                    if (a -> Bool
predicate a
x)
                    then s -> m a
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
val -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
val (forall a. a -> Maybe a
Just s
s)
                    else SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc'
                Skip s
s -> SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc
                Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
val -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
val forall a. Maybe a
Nothing

    stepOuter Fold m a a
_ State Stream m a
_ Maybe s
Nothing = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL groupsBy #-}
groupsBy :: Monad m
    => (a -> a -> Bool)
    -> Fold m a b
    -> Stream m a
    -> Stream m b
groupsBy :: forall (m :: * -> *) a b.
Monad m =>
(a -> a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
groupsBy a -> a -> Bool
cmp Fold m a b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (forall {a} {m :: * -> *} {a}.
Fold m a a
-> State Stream m a
-> (Maybe s, Maybe a)
-> m (Step (Maybe s, Maybe a) a)
stepOuter Fold m a b
f) (forall a. a -> Maybe a
Just s
state, forall a. Maybe a
Nothing)

    where

    {-# INLINE_LATE stepOuter #-}
    stepOuter :: Fold m a a
-> State Stream m a
-> (Maybe s, Maybe a)
-> m (Step (Maybe s, Maybe a) a)
stepOuter (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st, Maybe a
Nothing) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                s
acc' <- s -> a -> m s
fstep s
acc a
x
                SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
x s
s s
acc'

            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ (forall a. a -> Maybe a
Just s
s, forall a. Maybe a
Nothing)
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where

        go :: SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go !SPEC
_ a
prev s
stt !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
            case Step s a
res of
                Yield a
x s
s -> do
                    if a -> a -> Bool
cmp a
x a
prev
                    then do
                        s
acc' <- s -> a -> m s
fstep s
acc a
x
                        SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
prev s
s s
acc'
                    else s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. a -> Maybe a
Just s
s, forall a. a -> Maybe a
Just a
x)
                Skip s
s -> SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
prev s
s s
acc
                Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

    stepOuter (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st, Just a
prev) = do
        s
acc <- m s
initial
        s
acc' <- s -> a -> m s
fstep s
acc a
prev
        SPEC -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC s
st s
acc'

        where

        -- XXX code duplicated from the previous equation
        go :: SPEC -> s -> s -> m (Step (Maybe s, Maybe a) a)
go !SPEC
_ s
stt !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
            case Step s a
res of
                Yield a
x s
s -> do
                    if a -> a -> Bool
cmp a
x a
prev
                    then do
                        s
acc' <- s -> a -> m s
fstep s
acc a
x
                        SPEC -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC s
s s
acc'
                    else s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. a -> Maybe a
Just s
s, forall a. a -> Maybe a
Just a
x)
                Skip s
s -> SPEC -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC s
s s
acc
                Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

    stepOuter Fold m a a
_ State Stream m a
_ (Maybe s
Nothing,Maybe a
_) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL groupsRollingBy #-}
groupsRollingBy :: Monad m
    => (a -> a -> Bool)
    -> Fold m a b
    -> Stream m a
    -> Stream m b
groupsRollingBy :: forall (m :: * -> *) a b.
Monad m =>
(a -> a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
groupsRollingBy a -> a -> Bool
cmp Fold m a b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (forall {a} {m :: * -> *} {a}.
Fold m a a
-> State Stream m a
-> (Maybe s, Maybe a)
-> m (Step (Maybe s, Maybe a) a)
stepOuter Fold m a b
f) (forall a. a -> Maybe a
Just s
state, forall a. Maybe a
Nothing)
    where

      {-# INLINE_LATE stepOuter #-}
      stepOuter :: Fold m a a
-> State Stream m a
-> (Maybe s, Maybe a)
-> m (Step (Maybe s, Maybe a) a)
stepOuter (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st, Maybe a
Nothing) = do
          Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
          case Step s a
res of
              Yield a
x s
s -> do
                  s
acc <- m s
initial
                  s
acc' <- s -> a -> m s
fstep s
acc a
x
                  SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
x s
s s
acc'

              Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ (forall a. a -> Maybe a
Just s
s, forall a. Maybe a
Nothing)
              Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where
          go :: SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go !SPEC
_ a
prev s
stt !s
acc = do
              Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
              case Step s a
res of
                  Yield a
x s
s -> do
                      if a -> a -> Bool
cmp a
prev a
x
                        then do
                          s
acc' <- s -> a -> m s
fstep s
acc a
x
                          SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
x s
s s
acc'
                        else
                          s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. a -> Maybe a
Just s
s, forall a. a -> Maybe a
Just a
x)
                  Skip s
s -> SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
prev s
s s
acc
                  Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

      stepOuter (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st, Just a
prev') = do
          s
acc <- m s
initial
          s
acc' <- s -> a -> m s
fstep s
acc a
prev'
          SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
prev' s
st s
acc'

        where
          go :: SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go !SPEC
_ a
prevv s
stt !s
acc = do
              Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
              case Step s a
res of
                  Yield a
x s
s -> do
                      if a -> a -> Bool
cmp a
prevv a
x
                      then do
                          s
acc' <- s -> a -> m s
fstep s
acc a
x
                          SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
x s
s s
acc'
                      else s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. a -> Maybe a
Just s
s, forall a. a -> Maybe a
Just a
x)
                  Skip s
s -> SPEC -> a -> s -> s -> m (Step (Maybe s, Maybe a) a)
go SPEC
SPEC a
prevv s
s s
acc
                  Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

      stepOuter Fold m a a
_ State Stream m a
_ (Maybe s
Nothing, Maybe a
_) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL splitBy #-}
splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
splitBy :: forall (m :: * -> *) a b.
Monad m =>
(a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
splitBy a -> Bool
predicate Fold m a b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (forall {a} {m :: * -> *} {a}.
Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' Fold m a b
f) (forall a. a -> Maybe a
Just s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st) = m s
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
st

        where

        go :: SPEC -> s -> s -> m (Step (Maybe s) a)
go !SPEC
_ s
stt !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
            case Step s a
res of
                Yield a
x s
s -> do
                    if a -> Bool
predicate a
x
                    then s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. a -> Maybe a
Just s
s)
                    else do
                        s
acc' <- s -> a -> m s
fstep s
acc a
x
                        SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc'
                Skip s
s -> SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc
                Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r forall a. Maybe a
Nothing

    step' Fold m a a
_ State Stream m a
_ Maybe s
Nothing = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
{-# INLINE_NORMAL splitSuffixBy #-}
splitSuffixBy :: Monad m
    => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
splitSuffixBy :: forall (m :: * -> *) a b.
Monad m =>
(a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
splitSuffixBy a -> Bool
predicate Fold m a b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (forall {a} {m :: * -> *} {a}.
Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' Fold m a b
f) (forall a. a -> Maybe a
Just s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                if a -> Bool
predicate a
x
                then s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
val -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
val (forall a. a -> Maybe a
Just s
s)
                else do
                    s
acc' <- s -> a -> m s
fstep s
acc a
x
                    SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc'

            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where

        go :: SPEC -> s -> s -> m (Step (Maybe s) a)
go !SPEC
_ s
stt !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
            case Step s a
res of
                Yield a
x s
s -> do
                    if a -> Bool
predicate a
x
                    then s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. a -> Maybe a
Just s
s)
                    else do
                        s
acc' <- s -> a -> m s
fstep s
acc a
x
                        SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc'
                Skip s
s -> SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc
                Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r forall a. Maybe a
Nothing

    step' Fold m a a
_ State Stream m a
_ Maybe s
Nothing = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL wordsBy #-}
wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
wordsBy :: forall (m :: * -> *) a b.
Monad m =>
(a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
wordsBy a -> Bool
predicate Fold m a b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (forall {a} {m :: * -> *} {a}.
Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
stepOuter Fold m a b
f) (forall a. a -> Maybe a
Just s
state)

    where

    {-# INLINE_LATE stepOuter #-}
    stepOuter :: Fold m a a -> State Stream m a -> Maybe s -> m (Step (Maybe s) a)
stepOuter (Fold s -> a -> m s
fstep m s
initial s -> m a
done) State Stream m a
gst (Just s
st) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
res of
            Yield a
x s
s -> do
                if a -> Bool
predicate a
x
                then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
s)
                else do
                    s
acc <- m s
initial
                    s
acc' <- s -> a -> m s
fstep s
acc a
x
                    SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc'

            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where

        go :: SPEC -> s -> s -> m (Step (Maybe s) a)
go !SPEC
_ s
stt !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
            case Step s a
res of
                Yield a
x s
s -> do
                    if a -> Bool
predicate a
x
                    then s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r (forall a. a -> Maybe a
Just s
s)
                    else do
                        s
acc' <- s -> a -> m s
fstep s
acc a
x
                        SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc'
                Skip s
s -> SPEC -> s -> s -> m (Step (Maybe s) a)
go SPEC
SPEC s
s s
acc
                Step s a
Stop -> s -> m a
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
r forall a. Maybe a
Nothing

    stepOuter Fold m a a
_ State Stream m a
_ Maybe s
Nothing = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- String search algorithms:
-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html

{-
-- TODO can we unify the splitting operations using a splitting configuration
-- like in the split package.
--
data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)

data SplitOptions = SplitOptions
    { style    :: SplitStyle
    , withSep  :: Bool  -- ^ keep the separators in output
    -- , compact  :: Bool  -- ^ treat multiple consecutive separators as one
    -- , trimHead :: Bool  -- ^ drop blank at head
    -- , trimTail :: Bool  -- ^ drop blank at tail
    }
-}

data SplitOnState s a =
      GO_START
    | GO_EMPTY_PAT s
    | GO_SINGLE_PAT s a
    | GO_SHORT_PAT s
    | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
    | GO_DONE

{-# INLINE_NORMAL splitOn #-}
splitOn
    :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
    => Array a
    -> Fold m a b
    -> Stream m a
    -> Stream m b
splitOn :: forall (m :: * -> *) a b.
(MonadIO m, Storable a, Enum a, Eq a) =>
Array a -> Fold m a b -> Stream m a -> Stream m b
splitOn Array a
patArr (Fold s -> a -> m s
fstep m s
initial s -> m b
done) (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> SplitOnState s a -> m (Step (SplitOnState s a) b)
stepOuter forall s a. SplitOnState s a
GO_START

    where

    patLen :: Int
patLen = forall a. Storable a => Array a -> Int
A.length Array a
patArr
    maxIndex :: Int
maxIndex = Int
patLen forall a. Num a => a -> a -> a
- Int
1
    elemBits :: Int
elemBits = forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a) forall a. Num a => a -> a -> a
* Int
8

    {-# INLINE_LATE stepOuter #-}
    stepOuter :: State Stream m a
-> SplitOnState s a -> m (Step (SplitOnState s a) b)
stepOuter State Stream m a
_ SplitOnState s a
GO_START =
        if Int
patLen forall a. Eq a => a -> a -> Bool
== Int
0
        then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> SplitOnState s a
GO_EMPTY_PAT s
state
        else if Int
patLen forall a. Eq a => a -> a -> Bool
== Int
1
            then do
                a
r <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ (forall a. Storable a => Array a -> Int -> IO a
A.unsafeIndexIO Array a
patArr Int
0)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> a -> SplitOnState s a
GO_SINGLE_PAT s
state a
r
            else if forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a) forall a. Num a => a -> a -> a
* Int
patLen
                    forall a. Ord a => a -> a -> Bool
<= forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: Word)
                then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
state
                else do
                    (Ring a
rb, Ptr a
rhead) <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Int -> IO (Ring a, Ptr a)
RB.new Int
patLen
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> Ring a -> Ptr a -> SplitOnState s a
GO_KARP_RABIN s
state Ring a
rb Ptr a
rhead

    stepOuter State Stream m a
gst (GO_SINGLE_PAT s
stt a
pat) = m s
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go SPEC
SPEC s
stt

        where

        go :: SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go !SPEC
_ s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    if a
pat forall a. Eq a => a -> a -> Bool
== a
x
                    then do
                        b
r <- s -> m b
done s
acc
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> a -> SplitOnState s a
GO_SINGLE_PAT s
s a
pat)
                    else s -> a -> m s
fstep s
acc a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go SPEC
SPEC s
s
                Skip s
s -> SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go SPEC
SPEC s
s s
acc
                Step s a
Stop -> s -> m b
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

    stepOuter State Stream m a
gst (GO_SHORT_PAT s
stt) = m s
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall {a}.
SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC Int
0 (Word
0 :: Word) s
stt

        where

        mask :: Word
        mask :: Word
mask = (Word
1 forall a. Bits a => a -> Int -> a
`shiftL` (Int
elemBits forall a. Num a => a -> a -> a
* Int
patLen)) forall a. Num a => a -> a -> a
- Word
1

        addToWord :: a -> a -> a
addToWord a
wrd a
a = (a
wrd forall a. Bits a => a -> Int -> a
`shiftL` Int
elemBits) forall a. Bits a => a -> a -> a
.|. forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Enum a => a -> Int
fromEnum a
a)

        patWord :: Word
        patWord :: Word
patWord = Word
mask forall a. Bits a => a -> a -> a
.&. forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
A.foldl' forall {a} {a}. (Bits a, Num a, Enum a) => a -> a -> a
addToWord Word
0 Array a
patArr

        go0 :: SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 !SPEC
_ !Int
idx Word
wrd s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    let wrd' :: Word
wrd' = forall {a} {a}. (Bits a, Num a, Enum a) => a -> a -> a
addToWord Word
wrd a
x
                    if Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex
                    then do
                        if Word
wrd' forall a. Bits a => a -> a -> a
.&. Word
mask forall a. Eq a => a -> a -> Bool
== Word
patWord
                        then do
                            b
r <- s -> m b
done s
acc
                            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
s)
                        else forall {a}. SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word
wrd' s
s s
acc
                    else SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Word
wrd' s
s s
acc
                Skip s
s -> SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC Int
idx Word
wrd s
s s
acc
                Step s a
Stop -> do
                    s
acc' <- if Int
idx forall a. Eq a => a -> a -> Bool
/= Int
0
                            then Word -> Int -> s -> m s
go2 Word
wrd Int
idx s
acc
                            else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                    s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        {-# INLINE go1 #-}
        go1 :: SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 !SPEC
_ Word
wrd s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    let wrd' :: Word
wrd' = forall {a} {a}. (Bits a, Num a, Enum a) => a -> a -> a
addToWord Word
wrd a
x
                        old :: Word
old = (Word
mask forall a. Bits a => a -> a -> a
.&. Word
wrd) forall a. Bits a => a -> Int -> a
`shiftR` (Int
elemBits forall a. Num a => a -> a -> a
* (Int
patLen forall a. Num a => a -> a -> a
- Int
1))
                    s
acc' <- s -> a -> m s
fstep s
acc (forall a. Enum a => Int -> a
toEnum forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
old)
                    if Word
wrd' forall a. Bits a => a -> a -> a
.&. Word
mask forall a. Eq a => a -> a -> Bool
== Word
patWord
                    then s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
s)
                    else SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word
wrd' s
s s
acc'
                Skip s
s -> SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word
wrd s
s s
acc
                Step s a
Stop -> do
                    s
acc' <- Word -> Int -> s -> m s
go2 Word
wrd Int
patLen s
acc
                    s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        go2 :: Word -> Int -> s -> m s
go2 !Word
wrd !Int
n !s
acc | Int
n forall a. Ord a => a -> a -> Bool
> Int
0 = do
            let old :: Word
old = (Word
mask forall a. Bits a => a -> a -> a
.&. Word
wrd) forall a. Bits a => a -> Int -> a
`shiftR` (Int
elemBits forall a. Num a => a -> a -> a
* (Int
n forall a. Num a => a -> a -> a
- Int
1))
            s -> a -> m s
fstep s
acc (forall a. Enum a => Int -> a
toEnum forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
old) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Word -> Int -> s -> m s
go2 Word
wrd (Int
n forall a. Num a => a -> a -> a
- Int
1)
        go2 Word
_ Int
_ s
acc = forall (m :: * -> *) a. Monad m => a -> m a
return s
acc

    stepOuter State Stream m a
gst (GO_KARP_RABIN s
stt Ring a
rb Ptr a
rhead) = do
        m s
initial forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC Int
0 Ptr a
rhead s
stt

        where

        k :: Word32
k = Word32
2891336453 :: Word32
        coeff :: Word32
coeff = Word32
k forall a b. (Num a, Integral b) => a -> b -> a
^ Int
patLen
        addCksum :: Word32 -> a -> Word32
addCksum Word32
cksum a
a = Word32
cksum forall a. Num a => a -> a -> a
* Word32
k forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Enum a => a -> Int
fromEnum a
a)
        deltaCksum :: Word32 -> a -> a -> Word32
deltaCksum Word32
cksum a
old a
new =
            forall {a}. Enum a => Word32 -> a -> Word32
addCksum Word32
cksum a
new forall a. Num a => a -> a -> a
- Word32
coeff forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Enum a => a -> Int
fromEnum a
old)

        -- XXX shall we use a random starting hash or 1 instead of 0?
        patHash :: Word32
patHash = forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
A.foldl' forall {a}. Enum a => Word32 -> a -> Word32
addCksum Word32
0 Array a
patArr

        -- rh == ringHead
        go0 :: SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 !SPEC
_ !Int
idx !Ptr a
rh s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    Ptr a
rh' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
x
                    if Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex
                    then do
                        let fold :: (b -> a -> b) -> b -> Ring a -> b
fold = forall a b.
Storable a =>
Ptr a -> (b -> a -> b) -> b -> Ring a -> b
RB.unsafeFoldRing (forall a. Ring a -> Ptr a
RB.ringBound Ring a
rb)
                        let !ringHash :: Word32
ringHash = forall {b}. (b -> a -> b) -> b -> Ring a -> b
fold forall {a}. Enum a => Word32 -> a -> Word32
addCksum Word32
0 Ring a
rb
                        if Word32
ringHash forall a. Eq a => a -> a -> Bool
== Word32
patHash
                        then SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go2 SPEC
SPEC Word32
ringHash Ptr a
rh' s
s s
acc
                        else SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
ringHash Ptr a
rh' s
s s
acc
                    else SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Ptr a
rh' s
s s
acc
                Skip s
s -> SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC Int
idx Ptr a
rh s
s s
acc
                Step s a
Stop -> do
                    !s
acc' <- if Int
idx forall a. Eq a => a -> a -> Bool
/= Int
0
                             then forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingM Ptr a
rh s -> a -> m s
fstep s
acc Ring a
rb
                             else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                    s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        -- XXX Theoretically this code can do 4 times faster if GHC generates
        -- optimal code. If we use just "(cksum' == patHash)" condition it goes
        -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
        -- the generated code changes drastically and becomes 4x slower. Need
        -- to investigate what is going on with GHC.
        {-# INLINE go1 #-}
        go1 :: SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 !SPEC
_ !Word32
cksum !Ptr a
rh s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    a
old <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek Ptr a
rh
                    let cksum' :: Word32
cksum' = forall {a} {a}. (Enum a, Enum a) => Word32 -> a -> a -> Word32
deltaCksum Word32
cksum a
old a
x
                    s
acc' <- s -> a -> m s
fstep s
acc a
old

                    if (Word32
cksum' forall a. Eq a => a -> a -> Bool
== Word32
patHash)
                    then do
                        Ptr a
rh' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
x)
                        SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go2 SPEC
SPEC Word32
cksum' Ptr a
rh' s
s s
acc'
                    else do
                        Ptr a
rh' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
x)
                        SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
cksum' Ptr a
rh' s
s s
acc'
                Skip s
s -> SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
cksum Ptr a
rh s
s s
acc
                Step s a
Stop -> do
                    s
acc' <- forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingFullM Ptr a
rh s -> a -> m s
fstep s
acc Ring a
rb
                    s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        go2 :: SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go2 !SPEC
_ !Word32
cksum' !Ptr a
rh' s
s !s
acc' = do
            if forall a. Ring a -> Ptr a -> Array a -> Bool
RB.unsafeEqArray Ring a
rb Ptr a
rh' Array a
patArr
            then do
                b
r <- s -> m b
done s
acc'
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> Ring a -> Ptr a -> SplitOnState s a
GO_KARP_RABIN s
s Ring a
rb Ptr a
rhead)
            else SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
cksum' Ptr a
rh' s
s s
acc'

    stepOuter State Stream m a
gst (GO_EMPTY_PAT s
st) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                s
acc' <- s -> a -> m s
fstep s
acc a
x
                s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> SplitOnState s a
GO_EMPTY_PAT s
s)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> SplitOnState s a
GO_EMPTY_PAT s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    stepOuter State Stream m a
_ SplitOnState s a
GO_DONE = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL splitSuffixOn #-}
splitSuffixOn
    :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
    => Bool
    -> Array a
    -> Fold m a b
    -> Stream m a
    -> Stream m b
splitSuffixOn :: forall (m :: * -> *) a b.
(MonadIO m, Storable a, Enum a, Eq a) =>
Bool -> Array a -> Fold m a b -> Stream m a -> Stream m b
splitSuffixOn Bool
withSep Array a
patArr (Fold s -> a -> m s
fstep m s
initial s -> m b
done)
                (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> SplitOnState s a -> m (Step (SplitOnState s a) b)
stepOuter forall s a. SplitOnState s a
GO_START

    where

    patLen :: Int
patLen = forall a. Storable a => Array a -> Int
A.length Array a
patArr
    maxIndex :: Int
maxIndex = Int
patLen forall a. Num a => a -> a -> a
- Int
1
    elemBits :: Int
elemBits = forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a) forall a. Num a => a -> a -> a
* Int
8

    {-# INLINE_LATE stepOuter #-}
    stepOuter :: State Stream m a
-> SplitOnState s a -> m (Step (SplitOnState s a) b)
stepOuter State Stream m a
_ SplitOnState s a
GO_START =
        if Int
patLen forall a. Eq a => a -> a -> Bool
== Int
0
        then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> SplitOnState s a
GO_EMPTY_PAT s
state
        else if Int
patLen forall a. Eq a => a -> a -> Bool
== Int
1
             then do
                a
r <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ (forall a. Storable a => Array a -> Int -> IO a
A.unsafeIndexIO Array a
patArr Int
0)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> a -> SplitOnState s a
GO_SINGLE_PAT s
state a
r
             else if forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: a) forall a. Num a => a -> a -> a
* Int
patLen
                    forall a. Ord a => a -> a -> Bool
<= forall a. Storable a => a -> Int
sizeOf (forall a. HasCallStack => a
undefined :: Word)
                  then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
state
                  else do
                    (Ring a
rb, Ptr a
rhead) <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Int -> IO (Ring a, Ptr a)
RB.new Int
patLen
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> Ring a -> Ptr a -> SplitOnState s a
GO_KARP_RABIN s
state Ring a
rb Ptr a
rhead

    stepOuter State Stream m a
gst (GO_SINGLE_PAT s
stt a
pat) = do
        -- This first part is the only difference between splitOn and
        -- splitSuffixOn.
        -- If the last element is a separator do not issue a blank segment.
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                if a
pat forall a. Eq a => a -> a -> Bool
== a
x
                then do
                    s
acc' <- if Bool
withSep then s -> a -> m s
fstep s
acc a
x else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                    s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> a -> SplitOnState s a
GO_SINGLE_PAT s
s a
pat)
                else s -> a -> m s
fstep s
acc a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go SPEC
SPEC s
s
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ (forall s a. s -> a -> SplitOnState s a
GO_SINGLE_PAT s
s a
pat)
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where

        -- This is identical for splitOn and splitSuffixOn
        go :: SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go !SPEC
_ s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    if a
pat forall a. Eq a => a -> a -> Bool
== a
x
                    then do
                        s
acc' <- if Bool
withSep then s -> a -> m s
fstep s
acc a
x else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                        b
r <- s -> m b
done s
acc'
                        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> a -> SplitOnState s a
GO_SINGLE_PAT s
s a
pat)
                    else s -> a -> m s
fstep s
acc a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go SPEC
SPEC s
s
                Skip s
s -> SPEC -> s -> s -> m (Step (SplitOnState s a) b)
go SPEC
SPEC s
s s
acc
                Step s a
Stop -> s -> m b
done s
acc forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

    stepOuter State Stream m a
gst (GO_SHORT_PAT s
stt) = do

        -- Call "initial" only if the stream yields an element, otherwise we
        -- may call "initial" but never yield anything. initial may produce a
        -- side effect, therefore we will end up doing and discard a side
        -- effect.

        let idx :: Int
idx = Int
0
        let wrd :: Word
wrd = Word
0
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                let wrd' :: Word
wrd' = forall {a} {a}. (Bits a, Num a, Enum a) => a -> a -> a
addToWord Word
wrd a
x
                s
acc' <- if Bool
withSep then s -> a -> m s
fstep s
acc a
x else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                if Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex
                then do
                    if Word
wrd' forall a. Bits a => a -> a -> a
.&. Word
mask forall a. Eq a => a -> a -> Bool
== Word
patWord
                    then s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
s)
                    else forall {a}.
SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Word
wrd' s
s s
acc'
                else forall {a}.
SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Word
wrd' s
s s
acc'
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where

        mask :: Word
        mask :: Word
mask = (Word
1 forall a. Bits a => a -> Int -> a
`shiftL` (Int
elemBits forall a. Num a => a -> a -> a
* Int
patLen)) forall a. Num a => a -> a -> a
- Word
1

        addToWord :: a -> a -> a
addToWord a
wrd a
a = (a
wrd forall a. Bits a => a -> Int -> a
`shiftL` Int
elemBits) forall a. Bits a => a -> a -> a
.|. forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Enum a => a -> Int
fromEnum a
a)

        patWord :: Word
        patWord :: Word
patWord = Word
mask forall a. Bits a => a -> a -> a
.&. forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
A.foldl' forall {a} {a}. (Bits a, Num a, Enum a) => a -> a -> a
addToWord Word
0 Array a
patArr

        go0 :: SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 !SPEC
_ !Int
idx Word
wrd s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    let wrd' :: Word
wrd' = forall {a} {a}. (Bits a, Num a, Enum a) => a -> a -> a
addToWord Word
wrd a
x
                    s
acc' <- if Bool
withSep then s -> a -> m s
fstep s
acc a
x else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                    if Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex
                    then do
                        if Word
wrd' forall a. Bits a => a -> a -> a
.&. Word
mask forall a. Eq a => a -> a -> Bool
== Word
patWord
                        then do
                            b
r <- s -> m b
done s
acc'
                            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
s)
                        else forall {a}. SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word
wrd' s
s s
acc'
                    else SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Word
wrd' s
s s
acc'
                Skip s
s -> SPEC -> Int -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC Int
idx Word
wrd s
s s
acc
                Step s a
Stop -> do
                    if (Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex) Bool -> Bool -> Bool
&& (Word
wrd forall a. Bits a => a -> a -> a
.&. Word
mask forall a. Eq a => a -> a -> Bool
== Word
patWord)
                    then forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
                    else do
                        s
acc' <- if Int
idx forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
withSep
                                then Word -> Int -> s -> m s
go2 Word
wrd Int
idx s
acc
                                else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                        s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        {-# INLINE go1 #-}
        go1 :: SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 !SPEC
_ Word
wrd s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    let wrd' :: Word
wrd' = forall {a} {a}. (Bits a, Num a, Enum a) => a -> a -> a
addToWord Word
wrd a
x
                        old :: Word
old = (Word
mask forall a. Bits a => a -> a -> a
.&. Word
wrd) forall a. Bits a => a -> Int -> a
`shiftR` (Int
elemBits forall a. Num a => a -> a -> a
* (Int
patLen forall a. Num a => a -> a -> a
- Int
1))
                    s
acc' <- if Bool
withSep
                            then s -> a -> m s
fstep s
acc a
x
                            else s -> a -> m s
fstep s
acc (forall a. Enum a => Int -> a
toEnum forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
old)
                    if Word
wrd' forall a. Bits a => a -> a -> a
.&. Word
mask forall a. Eq a => a -> a -> Bool
== Word
patWord
                    then s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> SplitOnState s a
GO_SHORT_PAT s
s)
                    else SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word
wrd' s
s s
acc'
                Skip s
s -> SPEC -> Word -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word
wrd s
s s
acc
                Step s a
Stop ->
                    -- If the last sequence is a separator do not issue a blank
                    -- segment.
                    if Word
wrd forall a. Bits a => a -> a -> a
.&. Word
mask forall a. Eq a => a -> a -> Bool
== Word
patWord
                    then forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
                    else do
                        s
acc' <- if Bool
withSep
                                then forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                                else Word -> Int -> s -> m s
go2 Word
wrd Int
patLen s
acc
                        s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        go2 :: Word -> Int -> s -> m s
go2 !Word
wrd !Int
n !s
acc | Int
n forall a. Ord a => a -> a -> Bool
> Int
0 = do
            let old :: Word
old = (Word
mask forall a. Bits a => a -> a -> a
.&. Word
wrd) forall a. Bits a => a -> Int -> a
`shiftR` (Int
elemBits forall a. Num a => a -> a -> a
* (Int
n forall a. Num a => a -> a -> a
- Int
1))
            s -> a -> m s
fstep s
acc (forall a. Enum a => Int -> a
toEnum forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
old) forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Word -> Int -> s -> m s
go2 Word
wrd (Int
n forall a. Num a => a -> a -> a
- Int
1)
        go2 Word
_ Int
_ s
acc = forall (m :: * -> *) a. Monad m => a -> m a
return s
acc

    stepOuter State Stream m a
gst (GO_KARP_RABIN s
stt Ring a
rb Ptr a
rhead) = do
        let idx :: Int
idx = Int
0
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
stt
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                s
acc' <- if Bool
withSep then s -> a -> m s
fstep s
acc a
x else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                Ptr a
rh' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rhead a
x)
                if Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex
                then do
                    let fold :: (b -> a -> b) -> b -> Ring a -> b
fold = forall a b.
Storable a =>
Ptr a -> (b -> a -> b) -> b -> Ring a -> b
RB.unsafeFoldRing (forall a. Ring a -> Ptr a
RB.ringBound Ring a
rb)
                    let !ringHash :: Word32
ringHash = forall {b}. (b -> a -> b) -> b -> Ring a -> b
fold forall {a}. Enum a => Word32 -> a -> Word32
addCksum Word32
0 Ring a
rb
                    if Word32
ringHash forall a. Eq a => a -> a -> Bool
== Word32
patHash
                    then SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go2 SPEC
SPEC Word32
ringHash Ptr a
rh' s
s s
acc'
                    else SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Ptr a
rh' s
s s
acc'
                else SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Ptr a
rh' s
s s
acc'
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> Ring a -> Ptr a -> SplitOnState s a
GO_KARP_RABIN s
s Ring a
rb Ptr a
rhead)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

        where

        k :: Word32
k = Word32
2891336453 :: Word32
        coeff :: Word32
coeff = Word32
k forall a b. (Num a, Integral b) => a -> b -> a
^ Int
patLen
        addCksum :: Word32 -> a -> Word32
addCksum Word32
cksum a
a = Word32
cksum forall a. Num a => a -> a -> a
* Word32
k forall a. Num a => a -> a -> a
+ forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Enum a => a -> Int
fromEnum a
a)
        deltaCksum :: Word32 -> a -> a -> Word32
deltaCksum Word32
cksum a
old a
new =
            forall {a}. Enum a => Word32 -> a -> Word32
addCksum Word32
cksum a
new forall a. Num a => a -> a -> a
- Word32
coeff forall a. Num a => a -> a -> a
* forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Enum a => a -> Int
fromEnum a
old)

        -- XXX shall we use a random starting hash or 1 instead of 0?
        patHash :: Word32
patHash = forall a b. Storable a => (b -> a -> b) -> b -> Array a -> b
A.foldl' forall {a}. Enum a => Word32 -> a -> Word32
addCksum Word32
0 Array a
patArr

        -- rh == ringHead
        go0 :: SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 !SPEC
_ !Int
idx !Ptr a
rh s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    s
acc' <- if Bool
withSep then s -> a -> m s
fstep s
acc a
x else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                    Ptr a
rh' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
x)
                    if Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex
                    then do
                        let fold :: (b -> a -> b) -> b -> Ring a -> b
fold = forall a b.
Storable a =>
Ptr a -> (b -> a -> b) -> b -> Ring a -> b
RB.unsafeFoldRing (forall a. Ring a -> Ptr a
RB.ringBound Ring a
rb)
                        let !ringHash :: Word32
ringHash = forall {b}. (b -> a -> b) -> b -> Ring a -> b
fold forall {a}. Enum a => Word32 -> a -> Word32
addCksum Word32
0 Ring a
rb
                        if Word32
ringHash forall a. Eq a => a -> a -> Bool
== Word32
patHash
                        then SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go2 SPEC
SPEC Word32
ringHash Ptr a
rh' s
s s
acc'
                        else SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
ringHash Ptr a
rh' s
s s
acc'
                    else SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC (Int
idx forall a. Num a => a -> a -> a
+ Int
1) Ptr a
rh' s
s s
acc'
                Skip s
s -> SPEC -> Int -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go0 SPEC
SPEC Int
idx Ptr a
rh s
s s
acc
                Step s a
Stop -> do
                    -- do not issue a blank segment when we end at pattern
                    if (Int
idx forall a. Eq a => a -> a -> Bool
== Int
maxIndex) Bool -> Bool -> Bool
&& forall a. Ring a -> Ptr a -> Array a -> Bool
RB.unsafeEqArray Ring a
rb Ptr a
rh Array a
patArr
                    then forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
                    else do
                        !s
acc' <- if Int
idx forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
withSep
                                 then forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingM Ptr a
rh s -> a -> m s
fstep s
acc Ring a
rb
                                 else forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                        s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        -- XXX Theoretically this code can do 4 times faster if GHC generates
        -- optimal code. If we use just "(cksum' == patHash)" condition it goes
        -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
        -- the generated code changes drastically and becomes 4x slower. Need
        -- to investigate what is going on with GHC.
        {-# INLINE go1 #-}
        go1 :: SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 !SPEC
_ !Word32
cksum !Ptr a
rh s
st !s
acc = do
            Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
            case Step s a
res of
                Yield a
x s
s -> do
                    a
old <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ptr a -> IO a
peek Ptr a
rh
                    let cksum' :: Word32
cksum' = forall {a} {a}. (Enum a, Enum a) => Word32 -> a -> a -> Word32
deltaCksum Word32
cksum a
old a
x
                    s
acc' <- if Bool
withSep
                            then s -> a -> m s
fstep s
acc a
x
                            else s -> a -> m s
fstep s
acc a
old

                    if (Word32
cksum' forall a. Eq a => a -> a -> Bool
== Word32
patHash)
                    then do
                        Ptr a
rh' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
x)
                        SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go2 SPEC
SPEC Word32
cksum' Ptr a
rh' s
s s
acc'
                    else do
                        Ptr a
rh' <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
x)
                        SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
cksum' Ptr a
rh' s
s s
acc'
                Skip s
s -> SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
cksum Ptr a
rh s
s s
acc
                Step s a
Stop -> do
                    if forall a. Ring a -> Ptr a -> Array a -> Bool
RB.unsafeEqArray Ring a
rb Ptr a
rh Array a
patArr
                    then forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
                    else do
                        s
acc' <- if Bool
withSep
                                then forall (m :: * -> *) a. Monad m => a -> m a
return s
acc
                                else forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingFullM Ptr a
rh s -> a -> m s
fstep s
acc Ring a
rb
                        s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r forall s a. SplitOnState s a
GO_DONE

        go2 :: SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go2 !SPEC
_ !Word32
cksum' !Ptr a
rh' s
s !s
acc' = do
            if forall a. Ring a -> Ptr a -> Array a -> Bool
RB.unsafeEqArray Ring a
rb Ptr a
rh' Array a
patArr
            then do
                b
r <- s -> m b
done s
acc'
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> Ring a -> Ptr a -> SplitOnState s a
GO_KARP_RABIN s
s Ring a
rb Ptr a
rhead)
            else SPEC -> Word32 -> Ptr a -> s -> s -> m (Step (SplitOnState s a) b)
go1 SPEC
SPEC Word32
cksum' Ptr a
rh' s
s s
acc'

    stepOuter State Stream m a
gst (GO_EMPTY_PAT s
st) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
res of
            Yield a
x s
s -> do
                s
acc <- m s
initial
                s
acc' <- s -> a -> m s
fstep s
acc a
x
                s -> m b
done s
acc' forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
r (forall s a. s -> SplitOnState s a
GO_EMPTY_PAT s
s)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> SplitOnState s a
GO_EMPTY_PAT s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    stepOuter State Stream m a
_ SplitOnState s a
GO_DONE = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

data SplitState s arr
    = SplitInitial s
    | SplitBuffering s arr
    | SplitSplitting s arr
    | SplitYielding arr (SplitState s arr)
    | SplitFinishing

-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
-- split using a splitBy like combinator. The Fold would consume upto the
-- separator and return any leftover which can then be fed to the next fold.
--
-- We can revisit this once we have partial folds/parsers.
--
-- | Performs infix separator style splitting.
{-# INLINE_NORMAL splitInnerBy #-}
splitInnerBy
    :: Monad m
    => (f a -> m (f a, Maybe (f a)))  -- splitter
    -> (f a -> f a -> m (f a))        -- joiner
    -> Stream m (f a)
    -> Stream m (f a)
splitInnerBy :: forall (m :: * -> *) (f :: * -> *) a.
Monad m =>
(f a -> m (f a, Maybe (f a)))
-> (f a -> f a -> m (f a)) -> Stream m (f a) -> Stream m (f a)
splitInnerBy f a -> m (f a, Maybe (f a))
splitter f a -> f a -> m (f a)
joiner (Stream State Stream m (f a) -> s -> m (Step s (f a))
step1 s
state1) =
    (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m (f a)
-> SplitState s (f a) -> m (Step (SplitState s (f a)) (f a))
step (forall s arr. s -> SplitState s arr
SplitInitial s
state1))

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m (f a)
-> SplitState s (f a) -> m (Step (SplitState s (f a)) (f a))
step State Stream m (f a)
gst (SplitInitial s
st) = do
        Step s (f a)
r <- State Stream m (f a) -> s -> m (Step s (f a))
step1 State Stream m (f a)
gst s
st
        case Step s (f a)
r of
            Yield f a
x s
s -> do
                (f a
x1, Maybe (f a)
mx2) <- f a -> m (f a, Maybe (f a))
splitter f a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (f a)
mx2 of
                    Maybe (f a)
Nothing -> forall s a. s -> Step s a
Skip (forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
s f a
x1)
                    Just f a
x2 -> forall s a. s -> Step s a
Skip (forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
x1 (forall s arr. s -> arr -> SplitState s arr
SplitSplitting s
s f a
x2))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s arr. s -> SplitState s arr
SplitInitial s
s)
            Step s (f a)
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

    step State Stream m (f a)
gst (SplitBuffering s
st f a
buf) = do
        Step s (f a)
r <- State Stream m (f a) -> s -> m (Step s (f a))
step1 State Stream m (f a)
gst s
st
        case Step s (f a)
r of
            Yield f a
x s
s -> do
                (f a
x1, Maybe (f a)
mx2) <- f a -> m (f a, Maybe (f a))
splitter f a
x
                f a
buf' <- f a -> f a -> m (f a)
joiner f a
buf f a
x1
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (f a)
mx2 of
                    Maybe (f a)
Nothing -> forall s a. s -> Step s a
Skip (forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
s f a
buf')
                    Just f a
x2 -> forall s a. s -> Step s a
Skip (forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
buf' (forall s arr. s -> arr -> SplitState s arr
SplitSplitting s
s f a
x2))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
s f a
buf)
            Step s (f a)
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
buf forall s arr. SplitState s arr
SplitFinishing)

    step State Stream m (f a)
_ (SplitSplitting s
st f a
buf) = do
        (f a
x1, Maybe (f a)
mx2) <- f a -> m (f a, Maybe (f a))
splitter f a
buf
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (f a)
mx2 of
                Maybe (f a)
Nothing -> forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
st f a
x1
                Just f a
x2 -> forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
x1 (forall s arr. s -> arr -> SplitState s arr
SplitSplitting s
st f a
x2)

    step State Stream m (f a)
_ (SplitYielding f a
x SplitState s (f a)
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield f a
x SplitState s (f a)
next
    step State Stream m (f a)
_ SplitState s (f a)
SplitFinishing = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

-- | Performs infix separator style splitting.
{-# INLINE_NORMAL splitInnerBySuffix #-}
splitInnerBySuffix
    :: (Monad m, Eq (f a), Monoid (f a))
    => (f a -> m (f a, Maybe (f a)))  -- splitter
    -> (f a -> f a -> m (f a))        -- joiner
    -> Stream m (f a)
    -> Stream m (f a)
splitInnerBySuffix :: forall (m :: * -> *) (f :: * -> *) a.
(Monad m, Eq (f a), Monoid (f a)) =>
(f a -> m (f a, Maybe (f a)))
-> (f a -> f a -> m (f a)) -> Stream m (f a) -> Stream m (f a)
splitInnerBySuffix f a -> m (f a, Maybe (f a))
splitter f a -> f a -> m (f a)
joiner (Stream State Stream m (f a) -> s -> m (Step s (f a))
step1 s
state1) =
    (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m (f a)
-> SplitState s (f a) -> m (Step (SplitState s (f a)) (f a))
step (forall s arr. s -> SplitState s arr
SplitInitial s
state1))

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m (f a)
-> SplitState s (f a) -> m (Step (SplitState s (f a)) (f a))
step State Stream m (f a)
gst (SplitInitial s
st) = do
        Step s (f a)
r <- State Stream m (f a) -> s -> m (Step s (f a))
step1 State Stream m (f a)
gst s
st
        case Step s (f a)
r of
            Yield f a
x s
s -> do
                (f a
x1, Maybe (f a)
mx2) <- f a -> m (f a, Maybe (f a))
splitter f a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (f a)
mx2 of
                    Maybe (f a)
Nothing -> forall s a. s -> Step s a
Skip (forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
s f a
x1)
                    Just f a
x2 -> forall s a. s -> Step s a
Skip (forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
x1 (forall s arr. s -> arr -> SplitState s arr
SplitSplitting s
s f a
x2))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s arr. s -> SplitState s arr
SplitInitial s
s)
            Step s (f a)
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

    step State Stream m (f a)
gst (SplitBuffering s
st f a
buf) = do
        Step s (f a)
r <- State Stream m (f a) -> s -> m (Step s (f a))
step1 State Stream m (f a)
gst s
st
        case Step s (f a)
r of
            Yield f a
x s
s -> do
                (f a
x1, Maybe (f a)
mx2) <- f a -> m (f a, Maybe (f a))
splitter f a
x
                f a
buf' <- f a -> f a -> m (f a)
joiner f a
buf f a
x1
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (f a)
mx2 of
                    Maybe (f a)
Nothing -> forall s a. s -> Step s a
Skip (forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
s f a
buf')
                    Just f a
x2 -> forall s a. s -> Step s a
Skip (forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
buf' (forall s arr. s -> arr -> SplitState s arr
SplitSplitting s
s f a
x2))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
s f a
buf)
            Step s (f a)
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
                if f a
buf forall a. Eq a => a -> a -> Bool
== forall a. Monoid a => a
mempty
                then forall s a. Step s a
Stop
                else forall s a. s -> Step s a
Skip (forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
buf forall s arr. SplitState s arr
SplitFinishing)

    step State Stream m (f a)
_ (SplitSplitting s
st f a
buf) = do
        (f a
x1, Maybe (f a)
mx2) <- f a -> m (f a, Maybe (f a))
splitter f a
buf
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Maybe (f a)
mx2 of
                Maybe (f a)
Nothing -> forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s arr. s -> arr -> SplitState s arr
SplitBuffering s
st f a
x1
                Just f a
x2 -> forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s arr. arr -> SplitState s arr -> SplitState s arr
SplitYielding f a
x1 (forall s arr. s -> arr -> SplitState s arr
SplitSplitting s
st f a
x2)

    step State Stream m (f a)
_ (SplitYielding f a
x SplitState s (f a)
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield f a
x SplitState s (f a)
next
    step State Stream m (f a)
_ SplitState s (f a)
SplitFinishing = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Substreams
------------------------------------------------------------------------------

{-# INLINE_NORMAL isPrefixOf #-}
isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
isPrefixOf :: forall a (m :: * -> *).
(Eq a, Monad m) =>
Stream m a -> Stream m a -> m Bool
isPrefixOf (Stream State Stream m a -> s -> m (Step s a)
stepa s
ta) (Stream State Stream m a -> s -> m (Step s a)
stepb s
tb) = (s, s, Maybe a) -> m Bool
go (s
ta, s
tb, forall a. Maybe a
Nothing)
  where
    go :: (s, s, Maybe a) -> m Bool
go (s
sa, s
sb, Maybe a
Nothing) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepa forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
sa
        case Step s a
r of
            Yield a
x s
sa' -> (s, s, Maybe a) -> m Bool
go (s
sa', s
sb, forall a. a -> Maybe a
Just a
x)
            Skip s
sa'    -> (s, s, Maybe a) -> m Bool
go (s
sa', s
sb, forall a. Maybe a
Nothing)
            Step s a
Stop        -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

    go (s
sa, s
sb, Just a
x) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepb forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
sb
        case Step s a
r of
            Yield a
y s
sb' ->
                if a
x forall a. Eq a => a -> a -> Bool
== a
y
                    then (s, s, Maybe a) -> m Bool
go (s
sa, s
sb', forall a. Maybe a
Nothing)
                    else forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
            Skip s
sb' -> (s, s, Maybe a) -> m Bool
go (s
sa, s
sb', forall a. a -> Maybe a
Just a
x)
            Step s a
Stop     -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

{-# INLINE_NORMAL isSubsequenceOf #-}
isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
isSubsequenceOf :: forall a (m :: * -> *).
(Eq a, Monad m) =>
Stream m a -> Stream m a -> m Bool
isSubsequenceOf (Stream State Stream m a -> s -> m (Step s a)
stepa s
ta) (Stream State Stream m a -> s -> m (Step s a)
stepb s
tb) = (s, s, Maybe a) -> m Bool
go (s
ta, s
tb, forall a. Maybe a
Nothing)
  where
    go :: (s, s, Maybe a) -> m Bool
go (s
sa, s
sb, Maybe a
Nothing) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepa forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
sa
        case Step s a
r of
            Yield a
x s
sa' -> (s, s, Maybe a) -> m Bool
go (s
sa', s
sb, forall a. a -> Maybe a
Just a
x)
            Skip s
sa'    -> (s, s, Maybe a) -> m Bool
go (s
sa', s
sb, forall a. Maybe a
Nothing)
            Step s a
Stop        -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

    go (s
sa, s
sb, Just a
x) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepb forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
sb
        case Step s a
r of
            Yield a
y s
sb' ->
                if a
x forall a. Eq a => a -> a -> Bool
== a
y
                    then (s, s, Maybe a) -> m Bool
go (s
sa, s
sb', forall a. Maybe a
Nothing)
                    else (s, s, Maybe a) -> m Bool
go (s
sa, s
sb', forall a. a -> Maybe a
Just a
x)
            Skip s
sb' -> (s, s, Maybe a) -> m Bool
go (s
sa, s
sb', forall a. a -> Maybe a
Just a
x)
            Step s a
Stop     -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

{-# INLINE_NORMAL stripPrefix #-}
stripPrefix
    :: (Eq a, Monad m)
    => Stream m a -> Stream m a -> m (Maybe (Stream m a))
stripPrefix :: forall a (m :: * -> *).
(Eq a, Monad m) =>
Stream m a -> Stream m a -> m (Maybe (Stream m a))
stripPrefix (Stream State Stream m a -> s -> m (Step s a)
stepa s
ta) (Stream State Stream m a -> s -> m (Step s a)
stepb s
tb) = (s, s, Maybe a) -> m (Maybe (Stream m a))
go (s
ta, s
tb, forall a. Maybe a
Nothing)
  where
    go :: (s, s, Maybe a) -> m (Maybe (Stream m a))
go (s
sa, s
sb, Maybe a
Nothing) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepa forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
sa
        case Step s a
r of
            Yield a
x s
sa' -> (s, s, Maybe a) -> m (Maybe (Stream m a))
go (s
sa', s
sb, forall a. a -> Maybe a
Just a
x)
            Skip s
sa'    -> (s, s, Maybe a) -> m (Maybe (Stream m a))
go (s
sa', s
sb, forall a. Maybe a
Nothing)
            Step s a
Stop        -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
stepb s
sb)

    go (s
sa, s
sb, Just a
x) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepb forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
sb
        case Step s a
r of
            Yield a
y s
sb' ->
                if a
x forall a. Eq a => a -> a -> Bool
== a
y
                    then (s, s, Maybe a) -> m (Maybe (Stream m a))
go (s
sa, s
sb', forall a. Maybe a
Nothing)
                    else forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
            Skip s
sb' -> (s, s, Maybe a) -> m (Maybe (Stream m a))
go (s
sa, s
sb', forall a. a -> Maybe a
Just a
x)
            Step s a
Stop     -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

------------------------------------------------------------------------------
-- Map and Fold
------------------------------------------------------------------------------

-- | Execute a monadic action for each element of the 'Stream'
{-# INLINE_NORMAL mapM_ #-}
mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
mapM_ :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> m ()
mapM_ a -> m b
m = forall (m :: * -> *) a. Monad m => Stream m a -> m ()
drain forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
mapM a -> m b
m

-------------------------------------------------------------------------------
-- Stream transformations using Unfolds
-------------------------------------------------------------------------------

-- Define a unique structure to use in inspection testing
data ConcatMapUState o i =
      ConcatMapUOuter o
    | ConcatMapUInner o i

-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
-- to streams and then flattens the generated streams into a single output
-- stream.

-- This is like 'concatMap' but uses an unfold with an explicit state to
-- generate the stream instead of a 'Stream' type generator. This allows better
-- optimization via fusion.  This can be many times more efficient than
-- 'concatMap'.

{-# INLINE_NORMAL concatMapU #-}
concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
concatMapU :: forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> Stream m a -> Stream m b
concatMapU (Unfold s -> m (Step s b)
istep a -> m s
inject) (Stream State Stream m a -> s -> m (Step s a)
ostep s
ost) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> ConcatMapUState s s -> m (Step (ConcatMapUState s s) b)
step (forall o i. o -> ConcatMapUState o i
ConcatMapUOuter s
ost)
  where
    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> ConcatMapUState s s -> m (Step (ConcatMapUState s s) b)
step State Stream m a
gst (ConcatMapUOuter s
o) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
ostep (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
o
        case Step s a
r of
            Yield a
a s
o' -> do
                s
i <- a -> m s
inject a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall o i. o -> i -> ConcatMapUState o i
ConcatMapUInner s
o' s
i))
            Skip s
o' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. o -> ConcatMapUState o i
ConcatMapUOuter s
o')
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

    step State Stream m a
_ (ConcatMapUInner s
o s
i) = do
        Step s b
r <- s -> m (Step s b)
istep s
i
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
i' -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. o -> i -> ConcatMapUState o i
ConcatMapUInner s
o s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall o i. o -> i -> ConcatMapUState o i
ConcatMapUInner s
o s
i')
            Step s b
Stop       -> forall s a. s -> Step s a
Skip (forall o i. o -> ConcatMapUState o i
ConcatMapUOuter s
o)

data ConcatUnfoldInterleaveState o i =
      ConcatUnfoldInterleaveOuter o [i]
    | ConcatUnfoldInterleaveInner o [i]
    | ConcatUnfoldInterleaveInnerL [i] [i]
    | ConcatUnfoldInterleaveInnerR [i] [i]

-- XXX use arrays to store state instead of lists.
-- XXX In general we can use different scheduling strategies e.g. how to
-- schedule the outer vs inner loop or assigning weights to different streams
-- or outer and inner loops.

-- After a yield, switch to the next stream. Do not switch streams on Skip.
-- Yield from outer stream switches to the inner stream.
--
-- There are two choices here, (1) exhaust the outer stream first and then
-- start yielding from the inner streams, this is much simpler to implement,
-- (2) yield at least one element from an inner stream before going back to
-- outer stream and opening the next stream from it.
--
-- Ideally, we need some scheduling bias to inner streams vs outer stream.
-- Maybe we can configure the behavior.
--
{-# INLINE_NORMAL concatUnfoldInterleave #-}
concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
concatUnfoldInterleave :: forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> Stream m a -> Stream m b
concatUnfoldInterleave (Unfold s -> m (Step s b)
istep a -> m s
inject) (Stream State Stream m a -> s -> m (Step s a)
ostep s
ost) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> ConcatUnfoldInterleaveState s s
-> m (Step (ConcatUnfoldInterleaveState s s) b)
step (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
ost [])
  where
    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> ConcatUnfoldInterleaveState s s
-> m (Step (ConcatUnfoldInterleaveState s s) b)
step State Stream m a
gst (ConcatUnfoldInterleaveOuter s
o [s]
ls) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
ostep (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
o
        case Step s a
r of
            Yield a
a s
o' -> do
                s
i <- a -> m s
inject a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInner s
o' (s
i forall a. a -> [a] -> [a]
: [s]
ls)))
            Skip s
o' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
o' [s]
ls)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls [])

    step State Stream m a
_ (ConcatUnfoldInterleaveInner s
_ []) = forall a. HasCallStack => a
undefined
    step State Stream m a
_ (ConcatUnfoldInterleaveInner s
o (s
st:[s]
ls)) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
o (s
sforall a. a -> [a] -> [a]
:[s]
ls))
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInner s
o (s
sforall a. a -> [a] -> [a]
:[s]
ls))
            Step s b
Stop      -> forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
o [s]
ls)

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerL [] []) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step State Stream m a
_ (ConcatUnfoldInterleaveInnerL [] [s]
rs) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR [] [s]
rs)

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerL (s
st:[s]
ls) [s]
rs) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls (s
sforall a. a -> [a] -> [a]
:[s]
rs))
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL (s
sforall a. a -> [a] -> [a]
:[s]
ls) [s]
rs)
            Step s b
Stop      -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls [s]
rs)

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerR [] []) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step State Stream m a
_ (ConcatUnfoldInterleaveInnerR [s]
ls []) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls [])

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerR [s]
ls (s
st:[s]
rs)) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR (s
sforall a. a -> [a] -> [a]
:[s]
ls) [s]
rs)
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR [s]
ls (s
sforall a. a -> [a] -> [a]
:[s]
rs))
            Step s b
Stop      -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR [s]
ls [s]
rs)

-- XXX In general we can use different scheduling strategies e.g. how to
-- schedule the outer vs inner loop or assigning weights to different streams
-- or outer and inner loops.
--
-- This could be inefficient if the tasks are too small.
--
-- Compared to concatUnfoldInterleave this one switches streams on Skips.
--
{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
concatUnfoldRoundrobin :: forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> Stream m a -> Stream m b
concatUnfoldRoundrobin (Unfold s -> m (Step s b)
istep a -> m s
inject) (Stream State Stream m a -> s -> m (Step s a)
ostep s
ost) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> ConcatUnfoldInterleaveState s s
-> m (Step (ConcatUnfoldInterleaveState s s) b)
step (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
ost [])
  where
    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> ConcatUnfoldInterleaveState s s
-> m (Step (ConcatUnfoldInterleaveState s s) b)
step State Stream m a
gst (ConcatUnfoldInterleaveOuter s
o [s]
ls) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
ostep (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
o
        case Step s a
r of
            Yield a
a s
o' -> do
                s
i <- a -> m s
inject a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInner s
o' (s
i forall a. a -> [a] -> [a]
: [s]
ls)))
            Skip s
o' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInner s
o' [s]
ls)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls [])

    step State Stream m a
_ (ConcatUnfoldInterleaveInner s
o []) =
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
o [])

    step State Stream m a
_ (ConcatUnfoldInterleaveInner s
o (s
st:[s]
ls)) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
o (s
sforall a. a -> [a] -> [a]
:[s]
ls))
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
o (s
sforall a. a -> [a] -> [a]
:[s]
ls))
            Step s b
Stop      -> forall s a. s -> Step s a
Skip (forall o i. o -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveOuter s
o [s]
ls)

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerL [] []) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step State Stream m a
_ (ConcatUnfoldInterleaveInnerL [] [s]
rs) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR [] [s]
rs)

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerL (s
st:[s]
ls) [s]
rs) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls (s
sforall a. a -> [a] -> [a]
:[s]
rs))
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls (s
sforall a. a -> [a] -> [a]
:[s]
rs))
            Step s b
Stop      -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls [s]
rs)

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerR [] []) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step State Stream m a
_ (ConcatUnfoldInterleaveInnerR [s]
ls []) =
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerL [s]
ls [])

    step State Stream m a
_ (ConcatUnfoldInterleaveInnerR [s]
ls (s
st:[s]
rs)) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR (s
sforall a. a -> [a] -> [a]
:[s]
ls) [s]
rs)
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR (s
sforall a. a -> [a] -> [a]
:[s]
ls) [s]
rs)
            Step s b
Stop      -> forall s a. s -> Step s a
Skip (forall o i. [i] -> [i] -> ConcatUnfoldInterleaveState o i
ConcatUnfoldInterleaveInnerR [s]
ls [s]
rs)

data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2

-- Note that this could be much faster compared to the CPS stream. However, as
-- the number of streams being composed increases this may become expensive.
-- Need to see where the breaking point is between the two.
--
{-# INLINE_NORMAL append #-}
append :: Monad m => Stream m a -> Stream m a -> Stream m a
append :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> Stream m a -> Stream m a
append (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) (Stream State Stream m a -> s -> m (Step s a)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> AppendState s s -> m (Step (AppendState s s) a)
step (forall s1 s2. s1 -> AppendState s1 s2
AppendFirst s
state1)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a -> AppendState s s -> m (Step (AppendState s s) a)
step State Stream m a
gst (AppendFirst s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> AppendState s1 s2
AppendFirst s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> AppendState s1 s2
AppendFirst s
s)
            Step s a
Stop -> forall s a. s -> Step s a
Skip (forall s1 s2. s2 -> AppendState s1 s2
AppendSecond s
state2)

    step State Stream m a
gst (AppendSecond s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s2 -> AppendState s1 s2
AppendSecond s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s2 -> AppendState s1 s2
AppendSecond s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
    | InterleaveSecondOnly s2 | InterleaveFirstOnly s1

{-# INLINE_NORMAL interleave #-}
interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
interleave :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> Stream m a -> Stream m a
interleave (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) (Stream State Stream m a -> s -> m (Step s a)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
state1 s
state2)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step State Stream m a
gst (InterleaveFirst s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
s s
st2)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
s s
st2)
            Step s a
Stop -> forall s a. s -> Step s a
Skip (forall s1 s2. s2 -> InterleaveState s1 s2
InterleaveSecondOnly s
st2)

    step State Stream m a
gst (InterleaveSecond s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
st1 s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
st1 s
s)
            Step s a
Stop -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
st1)

    step State Stream m a
gst (InterleaveFirstOnly s
st1) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
gst (InterleaveSecondOnly s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s2 -> InterleaveState s1 s2
InterleaveSecondOnly s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s2 -> InterleaveState s1 s2
InterleaveSecondOnly s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

{-# INLINE_NORMAL interleaveMin #-}
interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
interleaveMin :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> Stream m a -> Stream m a
interleaveMin (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) (Stream State Stream m a -> s -> m (Step s a)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
state1 s
state2)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step State Stream m a
gst (InterleaveFirst s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
s s
st2)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
s s
st2)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
gst (InterleaveSecond s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
st1 s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
st1 s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
_ (InterleaveFirstOnly s
_) =  forall a. HasCallStack => a
undefined
    step State Stream m a
_ (InterleaveSecondOnly s
_) =  forall a. HasCallStack => a
undefined

{-# INLINE_NORMAL interleaveSuffix #-}
interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
interleaveSuffix :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> Stream m a -> Stream m a
interleaveSuffix (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) (Stream State Stream m a -> s -> m (Step s a)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
state1 s
state2)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step State Stream m a
gst (InterleaveFirst s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
s s
st2)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
s s
st2)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
gst (InterleaveSecond s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
st1 s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
st1 s
s)
            Step s a
Stop -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
st1)

    step State Stream m a
gst (InterleaveFirstOnly s
st1) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
_ (InterleaveSecondOnly s
_) =  forall a. HasCallStack => a
undefined

data InterleaveInfixState s1 s2 a
    = InterleaveInfixFirst s1 s2
    | InterleaveInfixSecondBuf s1 s2
    | InterleaveInfixSecondYield s1 s2 a
    | InterleaveInfixFirstYield s1 s2 a
    | InterleaveInfixFirstOnly s1

{-# INLINE_NORMAL interleaveInfix #-}
interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
interleaveInfix :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> Stream m a -> Stream m a
interleaveInfix (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) (Stream State Stream m a -> s -> m (Step s a)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> InterleaveInfixState s s a
-> m (Step (InterleaveInfixState s s a) a)
step (forall s1 s2 a. s1 -> s2 -> InterleaveInfixState s1 s2 a
InterleaveInfixFirst s
state1 s
state2)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> InterleaveInfixState s s a
-> m (Step (InterleaveInfixState s s a) a)
step State Stream m a
gst (InterleaveInfixFirst s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2 a. s1 -> s2 -> InterleaveInfixState s1 s2 a
InterleaveInfixSecondBuf s
s s
st2)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2 a. s1 -> s2 -> InterleaveInfixState s1 s2 a
InterleaveInfixFirst s
s s
st2)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
gst (InterleaveInfixSecondBuf s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. s -> Step s a
Skip (forall s1 s2 a. s1 -> s2 -> a -> InterleaveInfixState s1 s2 a
InterleaveInfixSecondYield s
st1 s
s a
a)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2 a. s1 -> s2 -> InterleaveInfixState s1 s2 a
InterleaveInfixSecondBuf s
st1 s
s)
            Step s a
Stop -> forall s a. s -> Step s a
Skip (forall s1 s2 a. s1 -> InterleaveInfixState s1 s2 a
InterleaveInfixFirstOnly s
st1)

    step State Stream m a
gst (InterleaveInfixSecondYield s
st1 s
st2 a
x) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
x (forall s1 s2 a. s1 -> s2 -> a -> InterleaveInfixState s1 s2 a
InterleaveInfixFirstYield s
s s
st2 a
a)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2 a. s1 -> s2 -> a -> InterleaveInfixState s1 s2 a
InterleaveInfixSecondYield s
s s
st2 a
x)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
_ (InterleaveInfixFirstYield s
st1 s
st2 a
x) = do
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall s1 s2 a. s1 -> s2 -> InterleaveInfixState s1 s2 a
InterleaveInfixSecondBuf s
st1 s
st2)

    step State Stream m a
gst (InterleaveInfixFirstOnly s
st1) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2 a. s1 -> InterleaveInfixState s1 s2 a
InterleaveInfixFirstOnly s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2 a. s1 -> InterleaveInfixState s1 s2 a
InterleaveInfixFirstOnly s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

{-# INLINE_NORMAL roundRobin #-}
roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
roundRobin :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> Stream m a -> Stream m a
roundRobin (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) (Stream State Stream m a -> s -> m (Step s a)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
state1 s
state2)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> InterleaveState s s -> m (Step (InterleaveState s s) a)
step State Stream m a
gst (InterleaveFirst s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
s s
st2)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveSecond s
s s
st2)
            Step s a
Stop -> forall s a. s -> Step s a
Skip (forall s1 s2. s2 -> InterleaveState s1 s2
InterleaveSecondOnly s
st2)

    step State Stream m a
gst (InterleaveSecond s
st1 s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
st1 s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> s2 -> InterleaveState s1 s2
InterleaveFirst s
st1 s
s)
            Step s a
Stop -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
st1)

    step State Stream m a
gst (InterleaveSecondOnly s
st2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step2 State Stream m a
gst s
st2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s2 -> InterleaveState s1 s2
InterleaveSecondOnly s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s2 -> InterleaveState s1 s2
InterleaveSecondOnly s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

    step State Stream m a
gst (InterleaveFirstOnly s
st1) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
s)
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> InterleaveState s1 s2
InterleaveFirstOnly s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

data ICUState s1 s2 i1 i2 =
      ICUFirst s1 s2
    | ICUSecond s1 s2
    | ICUSecondOnly s2
    | ICUFirstOnly s1
    | ICUFirstInner s1 s2 i1
    | ICUSecondInner s1 s2 i2
    | ICUFirstOnlyInner s1 i1
    | ICUSecondOnlyInner s2 i2

-- | Interleave streams (full streams, not the elements) unfolded from two
-- input streams and concat. Stop when the first stream stops. If the second
-- stream ends before the first one then first stream still keeps running alone
-- without any interleaving with the second stream.
--
--    [a1, a2, ... an]                   [b1, b2 ...]
-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
--
{-# INLINE_NORMAL gintercalateSuffix #-}
gintercalateSuffix
    :: Monad m
    => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
gintercalateSuffix :: forall (m :: * -> *) a c b.
Monad m =>
Unfold m a c
-> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
gintercalateSuffix
    (Unfold s -> m (Step s c)
istep1 a -> m s
inject1) (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1)
    (Unfold s -> m (Step s c)
istep2 b -> m s
inject2) (Stream State Stream m b -> s -> m (Step s b)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> ICUState s s s s -> m (Step (ICUState s s s s) c)
step (forall s1 s2 i1 i2. s1 -> s2 -> ICUState s1 s2 i1 i2
ICUFirst s
state1 s
state2)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> ICUState s s s s -> m (Step (ICUState s s s s) c)
step State Stream m a
gst (ICUFirst s
s1 s
s2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s a
r of
            Yield a
a s
s -> do
                s
i <- a -> m s
inject1 a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> i1 -> ICUState s1 s2 i1 i2
ICUFirstInner s
s s
s2 s
i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> ICUState s1 s2 i1 i2
ICUFirst s
s s
s2)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step State Stream m a
gst (ICUFirstOnly s
s1) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s a
r of
            Yield a
a s
s -> do
                s
i <- a -> m s
inject1 a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> i1 -> ICUState s1 s2 i1 i2
ICUFirstOnlyInner s
s s
i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> ICUState s1 s2 i1 i2
ICUFirstOnly s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step State Stream m a
_ (ICUFirstInner s
s1 s
s2 s
i1) = do
        Step s c
r <- s -> m (Step s c)
istep1 s
i1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 s2 i1 i2. s1 -> s2 -> i1 -> ICUState s1 s2 i1 i2
ICUFirstInner s
s1 s
s2 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> i1 -> ICUState s1 s2 i1 i2
ICUFirstInner s
s1 s
s2 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> ICUState s1 s2 i1 i2
ICUSecond s
s1 s
s2)

    step State Stream m a
_ (ICUFirstOnlyInner s
s1 s
i1) = do
        Step s c
r <- s -> m (Step s c)
istep1 s
i1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 s2 i1 i2. s1 -> i1 -> ICUState s1 s2 i1 i2
ICUFirstOnlyInner s
s1 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> i1 -> ICUState s1 s2 i1 i2
ICUFirstOnlyInner s
s1 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> ICUState s1 s2 i1 i2
ICUFirstOnly s
s1)

    step State Stream m a
gst (ICUSecond s
s1 s
s2) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
step2 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s2
        case Step s b
r of
            Yield b
a s
s -> do
                s
i <- b -> m s
inject2 b
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> i2 -> ICUState s1 s2 i1 i2
ICUSecondInner s
s1 s
s s
i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> ICUState s1 s2 i1 i2
ICUSecond s
s1 s
s)
            Step s b
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> ICUState s1 s2 i1 i2
ICUFirstOnly s
s1)

    step State Stream m a
_ (ICUSecondInner s
s1 s
s2 s
i2) = do
        Step s c
r <- s -> m (Step s c)
istep2 s
i2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 s2 i1 i2. s1 -> s2 -> i2 -> ICUState s1 s2 i1 i2
ICUSecondInner s
s1 s
s2 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> i2 -> ICUState s1 s2 i1 i2
ICUSecondInner s
s1 s
s2 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2. s1 -> s2 -> ICUState s1 s2 i1 i2
ICUFirst s
s1 s
s2)

    step State Stream m a
_ (ICUSecondOnly s
_s2) = forall a. HasCallStack => a
undefined
    step State Stream m a
_ (ICUSecondOnlyInner s
_s2 s
_i2) = forall a. HasCallStack => a
undefined

data InterposeSuffixState s1 i1 =
      InterposeSuffixFirst s1
    -- | InterposeSuffixFirstYield s1 i1
    | InterposeSuffixFirstInner s1 i1
    | InterposeSuffixSecond s1

-- Note that if an unfolded layer turns out to be nil we still emit the
-- separator effect. An alternate behavior could be to emit the separator
-- effect only if at least one element has been yielded by the unfolding.
-- However, that becomes a bit complicated, so we have chosen the former
-- behvaior for now.
{-# INLINE_NORMAL interposeSuffix #-}
interposeSuffix
    :: Monad m
    => m c -> Unfold m b c -> Stream m b -> Stream m c
interposeSuffix :: forall (m :: * -> *) c b.
Monad m =>
m c -> Unfold m b c -> Stream m b -> Stream m c
interposeSuffix
    m c
action
    (Unfold s -> m (Step s c)
istep1 b -> m s
inject1) (Stream State Stream m b -> s -> m (Step s b)
step1 s
state1) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> InterposeSuffixState s s
-> m (Step (InterposeSuffixState s s) c)
step (forall s1 i1. s1 -> InterposeSuffixState s1 i1
InterposeSuffixFirst s
state1)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> InterposeSuffixState s s
-> m (Step (InterposeSuffixState s s) c)
step State Stream m a
gst (InterposeSuffixFirst s
s1) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s b
r of
            Yield b
a s
s -> do
                s
i <- b -> m s
inject1 b
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 i1. s1 -> i1 -> InterposeSuffixState s1 i1
InterposeSuffixFirstInner s
s s
i))
                -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 i1. s1 -> InterposeSuffixState s1 i1
InterposeSuffixFirst s
s)
            Step s b
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    {-
    step _ (InterposeSuffixFirstYield s1 i1) = do
        r <- istep1 i1
        return $ case r of
            Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
            Skip i'    -> Skip (InterposeSuffixFirstYield s1 i')
            Stop       -> Skip (InterposeSuffixFirst s1)
    -}

    step State Stream m a
_ (InterposeSuffixFirstInner s
s1 s
i1) = do
        Step s c
r <- s -> m (Step s c)
istep1 s
i1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 i1. s1 -> i1 -> InterposeSuffixState s1 i1
InterposeSuffixFirstInner s
s1 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 i1. s1 -> i1 -> InterposeSuffixState s1 i1
InterposeSuffixFirstInner s
s1 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 i1. s1 -> InterposeSuffixState s1 i1
InterposeSuffixSecond s
s1)

    step State Stream m a
_ (InterposeSuffixSecond s
s1) = do
        c
r <- m c
action
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield c
r (forall s1 i1. s1 -> InterposeSuffixState s1 i1
InterposeSuffixFirst s
s1)

data ICALState s1 s2 i1 i2 a =
      ICALFirst s1 s2
    -- | ICALFirstYield s1 s2 i1
    | ICALFirstInner s1 s2 i1
    | ICALFirstOnly s1
    | ICALFirstOnlyInner s1 i1
    | ICALSecondInject s1 s2
    | ICALFirstInject s1 s2 i2
    -- | ICALFirstBuf s1 s2 i1 i2
    | ICALSecondInner s1 s2 i1 i2
    -- -- | ICALSecondInner s1 s2 i1 i2 a
    -- -- | ICALFirstResume s1 s2 i1 i2 a

-- | Interleave streams (full streams, not the elements) unfolded from two
-- input streams and concat. Stop when the first stream stops. If the second
-- stream ends before the first one then first stream still keeps running alone
-- without any interleaving with the second stream.
--
--    [a1, a2, ... an]                   [b1, b2 ...]
-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
--
{-# INLINE_NORMAL gintercalate #-}
gintercalate
    :: Monad m
    => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
gintercalate :: forall (m :: * -> *) a c b.
Monad m =>
Unfold m a c
-> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
gintercalate
    (Unfold s -> m (Step s c)
istep1 a -> m s
inject1) (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1)
    (Unfold s -> m (Step s c)
istep2 b -> m s
inject2) (Stream State Stream m b -> s -> m (Step s b)
step2 s
state2) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a} {a} {a}.
State Stream m a
-> ICALState s s s s a -> m (Step (ICALState s s s s a) c)
step (forall s1 s2 i1 i2 a. s1 -> s2 -> ICALState s1 s2 i1 i2 a
ICALFirst s
state1 s
state2)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> ICALState s s s s a -> m (Step (ICALState s s s s a) c)
step State Stream m a
gst (ICALFirst s
s1 s
s2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s a
r of
            Yield a
a s
s -> do
                s
i <- a -> m s
inject1 a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> i1 -> ICALState s1 s2 i1 i2 a
ICALFirstInner s
s s
s2 s
i))
                -- i `seq` return (Skip (ICALFirstYield s s2 i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> ICALState s1 s2 i1 i2 a
ICALFirst s
s s
s2)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    {-
    step _ (ICALFirstYield s1 s2 i1) = do
        r <- istep1 i1
        return $ case r of
            Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
            Skip i'    -> Skip (ICALFirstYield s1 s2 i')
            Stop       -> Skip (ICALFirst s1 s2)
    -}

    step State Stream m a
_ (ICALFirstInner s
s1 s
s2 s
i1) = do
        Step s c
r <- s -> m (Step s c)
istep1 s
i1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 s2 i1 i2 a. s1 -> s2 -> i1 -> ICALState s1 s2 i1 i2 a
ICALFirstInner s
s1 s
s2 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> i1 -> ICALState s1 s2 i1 i2 a
ICALFirstInner s
s1 s
s2 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> ICALState s1 s2 i1 i2 a
ICALSecondInject s
s1 s
s2)

    step State Stream m a
gst (ICALFirstOnly s
s1) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s a
r of
            Yield a
a s
s -> do
                s
i <- a -> m s
inject1 a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> i1 -> ICALState s1 s2 i1 i2 a
ICALFirstOnlyInner s
s s
i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> ICALState s1 s2 i1 i2 a
ICALFirstOnly s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step State Stream m a
_ (ICALFirstOnlyInner s
s1 s
i1) = do
        Step s c
r <- s -> m (Step s c)
istep1 s
i1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 s2 i1 i2 a. s1 -> i1 -> ICALState s1 s2 i1 i2 a
ICALFirstOnlyInner s
s1 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> i1 -> ICALState s1 s2 i1 i2 a
ICALFirstOnlyInner s
s1 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> ICALState s1 s2 i1 i2 a
ICALFirstOnly s
s1)

    -- We inject the second stream even before checking if the first stream
    -- would yield any more elements. There is no clear choice whether we
    -- should do this before or after that. Doing it after may make the state
    -- machine a bit simpler though.
    step State Stream m a
gst (ICALSecondInject s
s1 s
s2) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
step2 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s2
        case Step s b
r of
            Yield b
a s
s -> do
                s
i <- b -> m s
inject2 b
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> i2 -> ICALState s1 s2 i1 i2 a
ICALFirstInject s
s1 s
s s
i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> ICALState s1 s2 i1 i2 a
ICALSecondInject s
s1 s
s)
            Step s b
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> ICALState s1 s2 i1 i2 a
ICALFirstOnly s
s1)

    step State Stream m a
gst (ICALFirstInject s
s1 s
s2 s
i2) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s a
r of
            Yield a
a s
s -> do
                s
i <- a -> m s
inject1 a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a.
s1 -> s2 -> i1 -> i2 -> ICALState s1 s2 i1 i2 a
ICALSecondInner s
s s
s2 s
i s
i2))
                -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> i2 -> ICALState s1 s2 i1 i2 a
ICALFirstInject s
s s
s2 s
i2)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    {-
    step _ (ICALFirstBuf s1 s2 i1 i2) = do
        r <- istep1 i1
        return $ case r of
            Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
            Skip i'    -> Skip (ICALFirstBuf s1 s2 i' i2)
            Stop       -> Stop

    step _ (ICALSecondInner s1 s2 i1 i2 v) = do
        r <- istep2 i2
        return $ case r of
            Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
            Skip i'    -> Skip (ICALSecondInner s1 s2 i1 i' v)
            Stop       -> Skip (ICALFirstResume s1 s2 i1 i2 v)
    -}

    step State Stream m a
_ (ICALSecondInner s
s1 s
s2 s
i1 s
i2) = do
        Step s c
r <- s -> m (Step s c)
istep2 s
i2
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 s2 i1 i2 a.
s1 -> s2 -> i1 -> i2 -> ICALState s1 s2 i1 i2 a
ICALSecondInner s
s1 s
s2 s
i1 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a.
s1 -> s2 -> i1 -> i2 -> ICALState s1 s2 i1 i2 a
ICALSecondInner s
s1 s
s2 s
i1 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 s2 i1 i2 a. s1 -> s2 -> i1 -> ICALState s1 s2 i1 i2 a
ICALFirstInner s
s1 s
s2 s
i1)
            -- Stop       -> Skip (ICALFirstResume s1 s2 i1 i2)

    {-
    step _ (ICALFirstResume s1 s2 i1 i2 x) = do
        return $ Yield x (ICALFirstInner s1 s2 i1 i2)
    -}

data InterposeState s1 i1 a =
      InterposeFirst s1
    -- | InterposeFirstYield s1 i1
    | InterposeFirstInner s1 i1
    | InterposeFirstInject s1
    -- | InterposeFirstBuf s1 i1
    | InterposeSecondYield s1 i1
    -- -- | InterposeSecondYield s1 i1 a
    -- -- | InterposeFirstResume s1 i1 a

-- Note that this only interposes the pure values, we may run many effects to
-- generate those values as some effects may not generate anything (Skip).
{-# INLINE_NORMAL interpose #-}
interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
interpose :: forall (m :: * -> *) c b.
Monad m =>
m c -> Unfold m b c -> Stream m b -> Stream m c
interpose
    m c
action
    (Unfold s -> m (Step s c)
istep1 b -> m s
inject1) (Stream State Stream m b -> s -> m (Step s b)
step1 s
state1) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a} {a} {a}.
State Stream m a
-> InterposeState s s a -> m (Step (InterposeState s s a) c)
step (forall s1 i1 a. s1 -> InterposeState s1 i1 a
InterposeFirst s
state1)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> InterposeState s s a -> m (Step (InterposeState s s a) c)
step State Stream m a
gst (InterposeFirst s
s1) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s b
r of
            Yield b
a s
s -> do
                s
i <- b -> m s
inject1 b
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 i1 a. s1 -> i1 -> InterposeState s1 i1 a
InterposeFirstInner s
s s
i))
                -- i `seq` return (Skip (InterposeFirstYield s i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 i1 a. s1 -> InterposeState s1 i1 a
InterposeFirst s
s)
            Step s b
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    {-
    step _ (InterposeFirstYield s1 i1) = do
        r <- istep1 i1
        return $ case r of
            Yield x i' -> Yield x (InterposeFirstInner s1 i')
            Skip i'    -> Skip (InterposeFirstYield s1 i')
            Stop       -> Skip (InterposeFirst s1)
    -}

    step State Stream m a
_ (InterposeFirstInner s
s1 s
i1) = do
        Step s c
r <- s -> m (Step s c)
istep1 s
i1
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
i' -> forall s a. a -> s -> Step s a
Yield c
x (forall s1 i1 a. s1 -> i1 -> InterposeState s1 i1 a
InterposeFirstInner s
s1 s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall s1 i1 a. s1 -> i1 -> InterposeState s1 i1 a
InterposeFirstInner s
s1 s
i')
            Step s c
Stop       -> forall s a. s -> Step s a
Skip (forall s1 i1 a. s1 -> InterposeState s1 i1 a
InterposeFirstInject s
s1)

    step State Stream m a
gst (InterposeFirstInject s
s1) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s b
r of
            Yield b
a s
s -> do
                s
i <- b -> m s
inject1 b
a
                -- i `seq` return (Skip (InterposeFirstBuf s i))
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s1 i1 a. s1 -> i1 -> InterposeState s1 i1 a
InterposeSecondYield s
s s
i))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 i1 a. s1 -> InterposeState s1 i1 a
InterposeFirstInject s
s)
            Step s b
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    {-
    step _ (InterposeFirstBuf s1 i1) = do
        r <- istep1 i1
        return $ case r of
            Yield x i' -> Skip (InterposeSecondYield s1 i' x)
            Skip i'    -> Skip (InterposeFirstBuf s1 i')
            Stop       -> Stop
    -}

    {-
    step _ (InterposeSecondYield s1 i1 v) = do
        r <- action
        return $ Yield r (InterposeFirstResume s1 i1 v)
    -}
    step State Stream m a
_ (InterposeSecondYield s
s1 s
i1) = do
        c
r <- m c
action
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield c
r (forall s1 i1 a. s1 -> i1 -> InterposeState s1 i1 a
InterposeFirstInner s
s1 s
i1)

    {-
    step _ (InterposeFirstResume s1 i1 v) = do
        return $ Yield v (InterposeFirstInner s1 i1)
    -}

------------------------------------------------------------------------------
-- Exceptions
------------------------------------------------------------------------------

data GbracketState s1 s2 v
    = GBracketInit
    | GBracketNormal s1 v
    | GBracketException s2

-- | The most general bracketing and exception combinator. All other
-- combinators can be expressed in terms of this combinator. This can also be
-- used for cases which are not covered by the standard combinators.
--
-- /Internal/
--
{-# INLINE_NORMAL gbracket #-}
gbracket
    :: Monad m
    => m c                                  -- ^ before
    -> (forall s. m s -> m (Either e s))    -- ^ try (exception handling)
    -> (c -> m d)                           -- ^ after, on normal stop
    -> (c -> e -> Stream m b)               -- ^ on exception
    -> (c -> Stream m b)                    -- ^ stream generator
    -> Stream m b
gbracket :: forall (m :: * -> *) c e d b.
Monad m =>
m c
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> (c -> e -> Stream m b)
-> (c -> Stream m b)
-> Stream m b
gbracket m c
bef forall s. m s -> m (Either e s)
exc c -> m d
aft c -> e -> Stream m b
fexc c -> Stream m b
fnormal =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b
-> GbracketState (Stream m b) (Stream m b) c
-> m (Step (GbracketState (Stream m b) (Stream m b) c) b)
step forall s1 s2 v. GbracketState s1 s2 v
GBracketInit

    where

    {-# INLINE_LATE step #-}
    step :: State Stream m b
-> GbracketState (Stream m b) (Stream m b) c
-> m (Step (GbracketState (Stream m b) (Stream m b) c) b)
step State Stream m b
_ GbracketState (Stream m b) (Stream m b) c
GBracketInit = do
        c
r <- m c
bef
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s1 s2 v. s1 -> v -> GbracketState s1 s2 v
GBracketNormal (c -> Stream m b
fnormal c
r) c
r

    step State Stream m b
gst (GBracketNormal (UnStream State Stream m b -> s -> m (Step s b)
step1 s
st) c
v) = do
        Either e (Step s b)
res <- forall s. m s -> m (Either e s)
exc forall a b. (a -> b) -> a -> b
$ State Stream m b -> s -> m (Step s b)
step1 State Stream m b
gst s
st
        case Either e (Step s b)
res of
            Right Step s b
r -> case Step s b
r of
                Yield b
x s
s ->
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
x (forall s1 s2 v. s1 -> v -> GbracketState s1 s2 v
GBracketNormal (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s) c
v)
                Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 v. s1 -> v -> GbracketState s1 s2 v
GBracketNormal (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s) c
v)
                Step s b
Stop -> c -> m d
aft c
v forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
            Left e
e -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 v. s2 -> GbracketState s1 s2 v
GBracketException (c -> e -> Stream m b
fexc c
v e
e))
    step State Stream m b
gst (GBracketException (UnStream State Stream m b -> s -> m (Step s b)
step1 s
st)) = do
        Step s b
res <- State Stream m b -> s -> m (Step s b)
step1 State Stream m b
gst s
st
        case Step s b
res of
            Yield b
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
x (forall s1 s2 v. s2 -> GbracketState s1 s2 v
GBracketException (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s))
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 v. s2 -> GbracketState s1 s2 v
GBracketException (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s))
            Step s b
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- | Create an IORef holding a finalizer that is called automatically when the
-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
-- value to deactivate the finalizer.
newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
    => m a -> m (IORef (Maybe (IO ())))
newFinalizedIORef :: forall (m :: * -> *) a.
(MonadIO m, MonadBaseControl IO m) =>
m a -> m (IORef (Maybe (IO ())))
newFinalizedIORef m a
finalizer = do
    RunInIO m
mrun <- forall (m :: * -> *). MonadBaseControl IO m => m (RunInIO m)
captureMonadState
    IORef (Maybe (IO ()))
ref <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ do
                StM m a
_ <- forall (m :: * -> *). RunInIO m -> forall b. m b -> IO (StM m b)
runInIO RunInIO m
mrun m a
finalizer
                forall (m :: * -> *) a. Monad m => a -> m a
return ()
    let finalizer1 :: IO ()
finalizer1 = do
            Maybe (IO ())
res <- forall a. IORef a -> IO a
readIORef IORef (Maybe (IO ()))
ref
            case Maybe (IO ())
res of
                Maybe (IO ())
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
                Just IO ()
f -> IO ()
f
    Weak (IORef (Maybe (IO ())))
_ <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef (Maybe (IO ()))
ref IO ()
finalizer1
    forall (m :: * -> *) a. Monad m => a -> m a
return IORef (Maybe (IO ()))
ref

-- | Run the finalizer stored in an IORef and deactivate it so that it is run
-- only once.
--
runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
runIORefFinalizer :: forall (m :: * -> *). MonadIO m => IORef (Maybe (IO ())) -> m ()
runIORefFinalizer IORef (Maybe (IO ()))
ref = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
    Maybe (IO ())
res <- forall a. IORef a -> IO a
readIORef IORef (Maybe (IO ()))
ref
    case Maybe (IO ())
res of
        Maybe (IO ())
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return ()
        Just IO ()
f -> forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IO ()))
ref forall a. Maybe a
Nothing forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO ()
f

-- | Deactivate the finalizer stored in an IORef without running it.
--
clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
clearIORefFinalizer :: forall (m :: * -> *). MonadIO m => IORef (Maybe (IO ())) -> m ()
clearIORefFinalizer IORef (Maybe (IO ()))
ref = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> a -> IO ()
writeIORef IORef (Maybe (IO ()))
ref forall a. Maybe a
Nothing

data GbracketIOState s1 s2 v wref
    = GBracketIOInit
    | GBracketIONormal s1 v wref
    | GBracketIOException s2

-- | Like gbracket but also uses a finalizer to make sure when the stream is
-- garbage collected we run the finalizing action. This requires a MonadIO and
-- MonadBaseControl IO constraint.
--
-- | The most general bracketing and exception combinator. All other
-- combinators can be expressed in terms of this combinator. This can also be
-- used for cases which are not covered by the standard combinators.
--
-- /Internal/
--
{-# INLINE_NORMAL gbracketIO #-}
gbracketIO
    :: (MonadIO m, MonadBaseControl IO m)
    => m c                                  -- ^ before
    -> (forall s. m s -> m (Either e s))    -- ^ try (exception handling)
    -> (c -> m d)                           -- ^ after, on normal stop or GC
    -> (c -> e -> Stream m b)               -- ^ on exception
    -> (c -> Stream m b)                    -- ^ stream generator
    -> Stream m b
gbracketIO :: forall (m :: * -> *) c e d b.
(MonadIO m, MonadBaseControl IO m) =>
m c
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> (c -> e -> Stream m b)
-> (c -> Stream m b)
-> Stream m b
gbracketIO m c
bef forall s. m s -> m (Either e s)
exc c -> m d
aft c -> e -> Stream m b
fexc c -> Stream m b
fnormal =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b
-> GbracketIOState
     (Stream m b) (Stream m b) c (IORef (Maybe (IO ())))
-> m (Step
        (GbracketIOState
           (Stream m b) (Stream m b) c (IORef (Maybe (IO ()))))
        b)
step forall s1 s2 v wref. GbracketIOState s1 s2 v wref
GBracketIOInit

    where

    -- If the stream is never evaluated the "aft" action will never be
    -- called. For that to occur we will need the user of this API to pass a
    -- weak pointer to us.
    {-# INLINE_LATE step #-}
    step :: State Stream m b
-> GbracketIOState
     (Stream m b) (Stream m b) c (IORef (Maybe (IO ())))
-> m (Step
        (GbracketIOState
           (Stream m b) (Stream m b) c (IORef (Maybe (IO ()))))
        b)
step State Stream m b
_ GbracketIOState (Stream m b) (Stream m b) c (IORef (Maybe (IO ())))
GBracketIOInit = do
        -- We mask asynchronous exceptions to make the execution
        -- of 'bef' and the registration of 'aft' atomic.
        -- A similar thing is done in the resourcet package: https://git.io/JvKV3
        -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
        (c
r, IORef (Maybe (IO ()))
ref) <- forall (b :: * -> *) (m :: * -> *) a c.
MonadBaseControl b m =>
(b (StM m a) -> b (StM m c)) -> m a -> m c
liftBaseOp_ forall a. IO a -> IO a
mask_ forall a b. (a -> b) -> a -> b
$ do
            c
r <- m c
bef
            IORef (Maybe (IO ()))
ref <- forall (m :: * -> *) a.
(MonadIO m, MonadBaseControl IO m) =>
m a -> m (IORef (Maybe (IO ())))
newFinalizedIORef (c -> m d
aft c
r)
            forall (m :: * -> *) a. Monad m => a -> m a
return (c
r, IORef (Maybe (IO ()))
ref)
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s1 s2 v wref.
s1 -> v -> wref -> GbracketIOState s1 s2 v wref
GBracketIONormal (c -> Stream m b
fnormal c
r) c
r IORef (Maybe (IO ()))
ref

    step State Stream m b
gst (GBracketIONormal (UnStream State Stream m b -> s -> m (Step s b)
step1 s
st) c
v IORef (Maybe (IO ()))
ref) = do
        Either e (Step s b)
res <- forall s. m s -> m (Either e s)
exc forall a b. (a -> b) -> a -> b
$ State Stream m b -> s -> m (Step s b)
step1 State Stream m b
gst s
st
        case Either e (Step s b)
res of
            Right Step s b
r -> case Step s b
r of
                Yield b
x s
s ->
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
x (forall s1 s2 v wref.
s1 -> v -> wref -> GbracketIOState s1 s2 v wref
GBracketIONormal (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s) c
v IORef (Maybe (IO ()))
ref)
                Skip s
s ->
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 v wref.
s1 -> v -> wref -> GbracketIOState s1 s2 v wref
GBracketIONormal (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s) c
v IORef (Maybe (IO ()))
ref)
                Step s b
Stop -> do
                    forall (m :: * -> *). MonadIO m => IORef (Maybe (IO ())) -> m ()
runIORefFinalizer IORef (Maybe (IO ()))
ref
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
            Left e
e -> do
                forall (m :: * -> *). MonadIO m => IORef (Maybe (IO ())) -> m ()
clearIORefFinalizer IORef (Maybe (IO ()))
ref
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 v wref. s2 -> GbracketIOState s1 s2 v wref
GBracketIOException (c -> e -> Stream m b
fexc c
v e
e))
    step State Stream m b
gst (GBracketIOException (UnStream State Stream m b -> s -> m (Step s b)
step1 s
st)) = do
        Step s b
res <- State Stream m b -> s -> m (Step s b)
step1 State Stream m b
gst s
st
        case Step s b
res of
            Yield b
x s
s ->
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
x (forall s1 s2 v wref. s2 -> GbracketIOState s1 s2 v wref
GBracketIOException (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s))
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2 v wref. s2 -> GbracketIOState s1 s2 v wref
GBracketIOException (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
step1 s
s))
            Step s b
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- | Run a side effect before the stream yields its first element.
{-# INLINE_NORMAL before #-}
before :: Monad m => m b -> Stream m a -> Stream m a
before :: forall (m :: * -> *) b a.
Monad m =>
m b -> Stream m a -> Stream m a
before m b
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' forall a. Maybe a
Nothing

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' State Stream m a
_ Maybe s
Nothing = m b
action forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
state))

    step' State Stream m a
gst (Just s
st) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
res of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just s
s)
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
s)
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- | Run a side effect whenever the stream stops normally.
{-# INLINE_NORMAL after #-}
after :: Monad m => m b -> Stream m a -> Stream m a
after :: forall (m :: * -> *) b a.
Monad m =>
m b -> Stream m a -> Stream m a
after m b
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step' s
state

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
res of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x s
s
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip s
s
            Step s a
Stop      -> m b
action forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL afterIO #-}
afterIO :: (MonadIO m, MonadBaseControl IO m)
    => m b -> Stream m a -> Stream m a
afterIO :: forall (m :: * -> *) b a.
(MonadIO m, MonadBaseControl IO m) =>
m b -> Stream m a -> Stream m a
afterIO m b
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> Maybe (s, IORef (Maybe (IO ())))
-> m (Step (Maybe (s, IORef (Maybe (IO ())))) a)
step' forall a. Maybe a
Nothing

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> Maybe (s, IORef (Maybe (IO ())))
-> m (Step (Maybe (s, IORef (Maybe (IO ())))) a)
step' State Stream m a
_ Maybe (s, IORef (Maybe (IO ())))
Nothing = do
        IORef (Maybe (IO ()))
ref <- forall (m :: * -> *) a.
(MonadIO m, MonadBaseControl IO m) =>
m a -> m (IORef (Maybe (IO ())))
newFinalizedIORef m b
action
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (s
state, IORef (Maybe (IO ()))
ref)
    step' State Stream m a
gst (Just (s
st, IORef (Maybe (IO ()))
ref)) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
res of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just (s
s, IORef (Maybe (IO ()))
ref))
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (s
s, IORef (Maybe (IO ()))
ref))
            Step s a
Stop      -> do
                forall (m :: * -> *). MonadIO m => IORef (Maybe (IO ())) -> m ()
runIORefFinalizer IORef (Maybe (IO ()))
ref
                forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- XXX These combinators are expensive due to the call to
-- onException/handle/try on each step. Therefore, when possible, they should
-- be called in an outer loop where we perform less iterations. For example, we
-- cannot call them on each iteration in a char stream, instead we can call
-- them when doing an IO on an array.
--
-- XXX For high performance error checks in busy streams we may need another
-- Error constructor in step.
--
-- | Run a side effect whenever the stream aborts due to an exception. The
-- exception is not caught, simply rethrown.
{-# INLINE_NORMAL onException #-}
onException :: MonadCatch m => m b -> Stream m a -> Stream m a
onException :: forall (m :: * -> *) b a.
MonadCatch m =>
m b -> Stream m a -> Stream m a
onException m b
action Stream m a
str =
    forall (m :: * -> *) c e d b.
Monad m =>
m c
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> (c -> e -> Stream m b)
-> (c -> Stream m b)
-> Stream m b
gbracket (forall (m :: * -> *) a. Monad m => a -> m a
return ()) forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try forall (m :: * -> *) a. Monad m => a -> m a
return
        (\()
_ (SomeException
e :: MC.SomeException) -> forall (m :: * -> *) b a. Monad m => m b -> Stream m a
nilM (m b
action forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MC.throwM SomeException
e))
        (\()
_ -> Stream m a
str)

{-# INLINE_NORMAL _onException #-}
_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
_onException :: forall (m :: * -> *) b a.
MonadCatch m =>
m b -> Stream m a -> Stream m a
_onException m b
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step' s
state

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st forall (m :: * -> *) a b. MonadCatch m => m a -> m b -> m a
`MC.onException` m b
action
        case Step s a
res of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x s
s
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- XXX bracket is like concatMap, it generates a stream and then flattens it.
-- Like concatMap it has 10x worse performance compared to linear fused
-- compositions.
--
-- | Run the first action before the stream starts and remember its output,
-- generate a stream using the output, run the second action providing the
-- remembered value as an argument whenever the stream ends normally or due to
-- an exception.
{-# INLINE_NORMAL bracket #-}
bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
bracket :: forall (m :: * -> *) b c a.
MonadCatch m =>
m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
bracket m b
bef b -> m c
aft b -> Stream m a
bet =
    forall (m :: * -> *) c e d b.
Monad m =>
m c
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> (c -> e -> Stream m b)
-> (c -> Stream m b)
-> Stream m b
gbracket m b
bef forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try b -> m c
aft
        (\b
a (SomeException
e :: SomeException) -> forall (m :: * -> *) b a. Monad m => m b -> Stream m a
nilM (b -> m c
aft b
a forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MC.throwM SomeException
e)) b -> Stream m a
bet

{-# INLINE_NORMAL bracketIO #-}
bracketIO :: (MonadAsync m, MonadCatch m)
    => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
bracketIO :: forall (m :: * -> *) b c a.
(MonadAsync m, MonadCatch m) =>
m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
bracketIO m b
bef b -> m c
aft b -> Stream m a
bet =
    forall (m :: * -> *) c e d b.
(MonadIO m, MonadBaseControl IO m) =>
m c
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> (c -> e -> Stream m b)
-> (c -> Stream m b)
-> Stream m b
gbracketIO m b
bef forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try b -> m c
aft
        (\b
a (SomeException
e :: SomeException) -> forall (m :: * -> *) b a. Monad m => m b -> Stream m a
nilM (b -> m c
aft b
a forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MC.throwM SomeException
e)) b -> Stream m a
bet

data BracketState s v = BracketInit | BracketRun s v

{-# INLINE_NORMAL _bracket #-}
_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
_bracket :: forall (m :: * -> *) b c a.
MonadCatch m =>
m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
_bracket m b
bef b -> m c
aft b -> Stream m a
bet = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> BracketState (Stream m a) b
-> m (Step (BracketState (Stream m a) b) a)
step' forall s v. BracketState s v
BracketInit

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> BracketState (Stream m a) b
-> m (Step (BracketState (Stream m a) b) a)
step' State Stream m a
_ BracketState (Stream m a) b
BracketInit = m b
bef forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
x -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall s v. s -> v -> BracketState s v
BracketRun (b -> Stream m a
bet b
x) b
x))

    -- NOTE: It is important to use UnStream instead of the Stream pattern
    -- here, otherwise we get huge perf degradation, see note in concatMap.
    step' State Stream m a
gst (BracketRun (UnStream State Stream m a -> s -> m (Step s a)
step s
state) b
v) = do
        -- res <- step gst state `MC.onException` aft v
        Either SomeException (Step s a)
res <- forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try forall a b. (a -> b) -> a -> b
$ State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
state
        case Either SomeException (Step s a)
res of
            Left (SomeException
e :: SomeException) -> b -> m c
aft b
v forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
MC.throwM SomeException
e forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
            Right Step s a
r -> case Step s a
r of
                Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall s v. s -> v -> BracketState s v
BracketRun (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step s
s) b
v)
                Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s v. s -> v -> BracketState s v
BracketRun (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step s
s) b
v)
                Step s a
Stop      -> b -> m c
aft b
v forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-- | Run a side effect whenever the stream stops normally or aborts due to an
-- exception.
{-# INLINE finally #-}
finally :: MonadCatch m => m b -> Stream m a -> Stream m a
-- finally action xs = after action $ onException action xs
finally :: forall (m :: * -> *) b a.
MonadCatch m =>
m b -> Stream m a -> Stream m a
finally m b
action Stream m a
xs = forall (m :: * -> *) b c a.
MonadCatch m =>
m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
bracket (forall (m :: * -> *) a. Monad m => a -> m a
return ()) (\()
_ -> m b
action) (forall a b. a -> b -> a
const Stream m a
xs)

{-# INLINE finallyIO #-}
finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
finallyIO :: forall (m :: * -> *) b a.
(MonadAsync m, MonadCatch m) =>
m b -> Stream m a -> Stream m a
finallyIO m b
action Stream m a
xs = forall (m :: * -> *) b c a.
(MonadAsync m, MonadCatch m) =>
m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
bracketIO (forall (m :: * -> *) a. Monad m => a -> m a
return ()) (\()
_ -> m b
action) (forall a b. a -> b -> a
const Stream m a
xs)

-- | When evaluating a stream if an exception occurs, stream evaluation aborts
-- and the specified exception handler is run with the exception as argument.
{-# INLINE_NORMAL handle #-}
handle :: (MonadCatch m, Exception e)
    => (e -> Stream m a) -> Stream m a -> Stream m a
handle :: forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
(e -> Stream m a) -> Stream m a -> Stream m a
handle e -> Stream m a
f Stream m a
str =
    forall (m :: * -> *) c e d b.
Monad m =>
m c
-> (forall s. m s -> m (Either e s))
-> (c -> m d)
-> (c -> e -> Stream m b)
-> (c -> Stream m b)
-> Stream m b
gbracket (forall (m :: * -> *) a. Monad m => a -> m a
return ()) forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try forall (m :: * -> *) a. Monad m => a -> m a
return (\()
_ e
e -> e -> Stream m a
f e
e) (\()
_ -> Stream m a
str)

{-# INLINE_NORMAL _handle #-}
_handle :: (MonadCatch m, Exception e)
    => (e -> Stream m a) -> Stream m a -> Stream m a
_handle :: forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
(e -> Stream m a) -> Stream m a -> Stream m a
_handle e -> Stream m a
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> Either s (Stream m a) -> m (Step (Either s (Stream m a)) a)
step' (forall a b. a -> Either a b
Left s
state)

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> Either s (Stream m a) -> m (Step (Either s (Stream m a)) a)
step' State Stream m a
gst (Left s
st) = do
        Either e (Step s a)
res <- forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> m (Either e a)
MC.try forall a b. (a -> b) -> a -> b
$ State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Either e (Step s a)
res of
            Left e
e -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right (e -> Stream m a
f e
e)
            Right Step s a
r -> case Step s a
r of
                Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a b. a -> Either a b
Left s
s)
                Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
s)
                Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
gst (Right (UnStream State Stream m a -> s -> m (Step s a)
step1 s
st)) = do
        Step s a
res <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        case Step s a
res of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a b. b -> Either a b
Right (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step1 s
s))
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step1 s
s))
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

-------------------------------------------------------------------------------
-- General transformation
-------------------------------------------------------------------------------

{-# INLINE_NORMAL transform #-}
transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
transform :: forall (m :: * -> *) a b.
Monad m =>
Pipe m a b -> Stream m a -> Stream m b
transform (Pipe s1 -> a -> m (Step (PipeState s1 s2) b)
pstep1 s2 -> m (Step (PipeState s1 s2) b)
pstep2 s1
pstate) (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> (PipeState s1 s2, s) -> m (Step (PipeState s1 s2, s) b)
step' (forall s1 s2. s1 -> PipeState s1 s2
Consume s1
pstate, s
state)

  where

    {-# INLINE_LATE step' #-}

    step' :: State Stream m a
-> (PipeState s1 s2, s) -> m (Step (PipeState s1 s2, s) b)
step' State Stream m a
gst (Consume s1
pst, s
st) = s1
pst seq :: forall a b. a -> b -> b
`seq` do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Step (PipeState s1 s2) b
res <- s1 -> a -> m (Step (PipeState s1 s2) b)
pstep1 s1
pst a
x
                case Step (PipeState s1 s2) b
res of
                    Pipe.Yield b
b PipeState s1 s2
pst' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
b (PipeState s1 s2
pst', s
s)
                    Pipe.Continue PipeState s1 s2
pst' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (PipeState s1 s2
pst', s
s)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s1 s2. s1 -> PipeState s1 s2
Consume s1
pst, s
s)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
_ (Produce s2
pst, s
st) = s2
pst seq :: forall a b. a -> b -> b
`seq` do
        Step (PipeState s1 s2) b
res <- s2 -> m (Step (PipeState s1 s2) b)
pstep2 s2
pst
        case Step (PipeState s1 s2) b
res of
            Pipe.Yield b
b PipeState s1 s2
pst' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
b (PipeState s1 s2
pst', s
st)
            Pipe.Continue PipeState s1 s2
pst' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (PipeState s1 s2
pst', s
st)

------------------------------------------------------------------------------
-- Transformation by Folding (Scans)
------------------------------------------------------------------------------

------------------------------------------------------------------------------
-- Prescans
------------------------------------------------------------------------------

-- XXX Is a prescan useful, discarding the last step does not sound useful?  I
-- am not sure about the utility of this function, so this is implemented but
-- not exposed. We can expose it if someone provides good reasons why this is
-- useful.
--
-- XXX We have to execute the stream one step ahead to know that we are at the
-- last step.  The vector implementation of prescan executes the last fold step
-- but does not yield the result. This means we have executed the effect but
-- discarded value. This does not sound right. In this implementation we are
-- not executing the last fold step.
{-# INLINE_NORMAL prescanlM' #-}
prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
prescanlM' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> Stream m b
prescanlM' b -> a -> m b
f m b
mz (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a -> (s, m b) -> m (Step (s, m b) b)
step' (s
state, m b
mz)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, m b) -> m (Step (s, m b) b)
step' State Stream m a
gst (s
st, m b
prev) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> do
                b
acc <- m b
prev
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
acc (s
s, b -> a -> m b
f b
acc a
x)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, m b
prev)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE prescanl' #-}
prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
prescanl' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> Stream m b
prescanl' b -> a -> b
f b
z = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> Stream m b
prescanlM' (\b
a a
b -> forall (m :: * -> *) a. Monad m => a -> m a
return (b -> a -> b
f b
a a
b)) (forall (m :: * -> *) a. Monad m => a -> m a
return b
z)

------------------------------------------------------------------------------
-- Monolithic postscans (postscan followed by a map)
------------------------------------------------------------------------------

-- The performance of a modular postscan followed by a map seems to be
-- equivalent to this monolithic scan followed by map therefore we may not need
-- this implementation. We just have it for performance comparison and in case
-- modular version does not perform well in some situation.
--
{-# INLINE_NORMAL postscanlMx' #-}
postscanlMx' :: Monad m
    => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
postscanlMx' :: forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
postscanlMx' x -> a -> m x
fstep m x
begin x -> m b
done (Stream State Stream m a -> s -> m (Step s a)
step s
state) = do
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a -> (s, m x) -> m (Step (s, m x) b)
step' (s
state, m x
begin)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, m x) -> m (Step (s, m x) b)
step' State Stream m a
gst (s
st, m x
acc) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> do
                x
old <- m x
acc
                x
y <- x -> a -> m x
fstep x
old a
x
                b
v <- x -> m b
done x
y
                b
v seq :: forall a b. a -> b -> b
`seq` x
y seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. a -> s -> Step s a
Yield b
v (s
s, forall (m :: * -> *) a. Monad m => a -> m a
return x
y))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, m x
acc)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL postscanlx' #-}
postscanlx' :: Monad m
    => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
postscanlx' :: forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
postscanlx' x -> a -> x
fstep x
begin x -> b
done Stream m a
s =
    forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
postscanlMx' (\x
b a
a -> forall (m :: * -> *) a. Monad m => a -> m a
return (x -> a -> x
fstep x
b a
a)) (forall (m :: * -> *) a. Monad m => a -> m a
return x
begin) (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> b
done) Stream m a
s

-- XXX do we need consM strict to evaluate the begin value?
{-# INLINE scanlMx' #-}
scanlMx' :: Monad m
    => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
scanlMx' :: forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
scanlMx' x -> a -> m x
fstep m x
begin x -> m b
done Stream m a
s =
    (m x
begin forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \x
x -> x
x seq :: forall a b. a -> b -> b
`seq` x -> m b
done x
x) forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
`consM` forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
postscanlMx' x -> a -> m x
fstep m x
begin x -> m b
done Stream m a
s

{-# INLINE scanlx' #-}
scanlx' :: Monad m
    => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
scanlx' :: forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
scanlx' x -> a -> x
fstep x
begin x -> b
done Stream m a
s =
    forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
scanlMx' (\x
b a
a -> forall (m :: * -> *) a. Monad m => a -> m a
return (x -> a -> x
fstep x
b a
a)) (forall (m :: * -> *) a. Monad m => a -> m a
return x
begin) (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> b
done) Stream m a
s

------------------------------------------------------------------------------
-- postscans
------------------------------------------------------------------------------

{-# INLINE_NORMAL postscanlM' #-}
postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM' b -> a -> m b
fstep b
begin (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    b
begin seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a -> (s, b) -> m (Step (s, b) b)
step' (s
state, b
begin)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, b) -> m (Step (s, b) b)
step' State Stream m a
gst (s
st, b
acc) = b
acc seq :: forall a b. a -> b -> b
`seq` do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> do
                b
y <- b -> a -> m b
fstep b
acc a
x
                b
y seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. a -> s -> Step s a
Yield b
y (s
s, b
y))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, b
acc)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL postscanl' #-}
postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
postscanl' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> Stream m b
postscanl' a -> b -> a
f = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM' (\a
a b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b -> a
f a
a b
b))

{-# INLINE_NORMAL postscanlM #-}
postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM b -> a -> m b
fstep b
begin (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a -> (s, b) -> m (Step (s, b) b)
step' (s
state, b
begin)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, b) -> m (Step (s, b) b)
step' State Stream m a
gst (s
st, b
acc) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> do
                b
y <- b -> a -> m b
fstep b
acc a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. a -> s -> Step s a
Yield b
y (s
s, b
y))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, b
acc)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL postscanl #-}
postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
postscanl :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> Stream m b
postscanl a -> b -> a
f = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM (\a
a b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b -> a
f a
a b
b))

{-# INLINE_NORMAL scanlM' #-}
scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
scanlM' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
scanlM' b -> a -> m b
fstep b
begin Stream m a
s = b
begin seq :: forall a b. a -> b -> b
`seq` (b
begin forall (m :: * -> *) a. Monad m => a -> Stream m a -> Stream m a
`cons` forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM' b -> a -> m b
fstep b
begin Stream m a
s)

{-# INLINE scanl' #-}
scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
scanl' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> Stream m b
scanl' b -> a -> b
f = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
scanlM' (\b
a a
b -> forall (m :: * -> *) a. Monad m => a -> m a
return (b -> a -> b
f b
a a
b))

{-# INLINE_NORMAL scanlM #-}
scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
scanlM :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
scanlM b -> a -> m b
fstep b
begin Stream m a
s = b
begin forall (m :: * -> *) a. Monad m => a -> Stream m a -> Stream m a
`cons` forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
postscanlM b -> a -> m b
fstep b
begin Stream m a
s

{-# INLINE scanl #-}
scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
scanl :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> Stream m b
scanl b -> a -> b
f = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> Stream m a -> Stream m b
scanlM (\b
a a
b -> forall (m :: * -> *) a. Monad m => a -> m a
return (b -> a -> b
f b
a a
b))

{-# INLINE_NORMAL scanl1M #-}
scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
scanl1M :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> m a) -> Stream m a -> Stream m a
scanl1M a -> a -> m a
fstep (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> (s, Maybe a) -> m (Step (s, Maybe a) a)
step' (s
state, forall a. Maybe a
Nothing)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, Maybe a) -> m (Step (s, Maybe a) a)
step' State Stream m a
gst (s
st, Maybe a
Nothing) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (s
s, forall a. a -> Maybe a
Just a
x)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, forall a. Maybe a
Nothing)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
gst (s
st, Just a
acc) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
y s
s -> do
                a
z <- a -> a -> m a
fstep a
acc a
y
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
z (s
s, forall a. a -> Maybe a
Just a
z)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, forall a. a -> Maybe a
Just a
acc)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE scanl1 #-}
scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
scanl1 :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> a) -> Stream m a -> Stream m a
scanl1 a -> a -> a
f = forall (m :: * -> *) a.
Monad m =>
(a -> a -> m a) -> Stream m a -> Stream m a
scanl1M (\a
x a
y -> forall (m :: * -> *) a. Monad m => a -> m a
return (a -> a -> a
f a
x a
y))

{-# INLINE_NORMAL scanl1M' #-}
scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
scanl1M' :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> m a) -> Stream m a -> Stream m a
scanl1M' a -> a -> m a
fstep (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> (s, Maybe a) -> m (Step (s, Maybe a) a)
step' (s
state, forall a. Maybe a
Nothing)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, Maybe a) -> m (Step (s, Maybe a) a)
step' State Stream m a
gst (s
st, Maybe a
Nothing) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> a
x seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (s
s, forall a. a -> Maybe a
Just a
x)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, forall a. Maybe a
Nothing)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
gst (s
st, Just a
acc) = a
acc seq :: forall a b. a -> b -> b
`seq` do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
y s
s -> do
                a
z <- a -> a -> m a
fstep a
acc a
y
                a
z seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
z (s
s, forall a. a -> Maybe a
Just a
z)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, forall a. a -> Maybe a
Just a
acc)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE scanl1' #-}
scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
scanl1' :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> a) -> Stream m a -> Stream m a
scanl1' a -> a -> a
f = forall (m :: * -> *) a.
Monad m =>
(a -> a -> m a) -> Stream m a -> Stream m a
scanl1M' (\a
x a
y -> forall (m :: * -> *) a. Monad m => a -> m a
return (a -> a -> a
f a
x a
y))

------------------------------------------------------------------------------
-- Stateful map/scan
------------------------------------------------------------------------------

data RollingMapState s a = RollingMapInit s | RollingMapGo s a

{-# INLINE rollingMapM #-}
rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
rollingMapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> a -> m b) -> Stream m a -> Stream m b
rollingMapM a -> a -> m b
f (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> RollingMapState s a -> m (Step (RollingMapState s a) b)
step (forall s a. s -> RollingMapState s a
RollingMapInit s
state1)
    where
    step :: State Stream m a
-> RollingMapState s a -> m (Step (RollingMapState s a) b)
step State Stream m a
gst (RollingMapInit s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> a -> RollingMapState s a
RollingMapGo s
s a
x
            Skip s
s -> forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> RollingMapState s a
RollingMapInit s
s
            Step s a
Stop   -> forall s a. Step s a
Stop

    step State Stream m a
gst (RollingMapGo s
s1 a
x1) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
s1
        case Step s a
r of
            Yield a
x s
s -> do
                !b
res <- a -> a -> m b
f a
x a
x1
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
res forall a b. (a -> b) -> a -> b
$ forall s a. s -> a -> RollingMapState s a
RollingMapGo s
s a
x
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall s a. s -> a -> RollingMapState s a
RollingMapGo s
s a
x1
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

{-# INLINE rollingMap #-}
rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
rollingMap :: forall (m :: * -> *) a b.
Monad m =>
(a -> a -> b) -> Stream m a -> Stream m b
rollingMap a -> a -> b
f = forall (m :: * -> *) a b.
Monad m =>
(a -> a -> m b) -> Stream m a -> Stream m b
rollingMapM (\a
x a
y -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ a -> a -> b
f a
x a
y)

------------------------------------------------------------------------------
-- Tapping/Distributing
------------------------------------------------------------------------------

{-# INLINE tap #-}
tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
tap :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> Stream m a
tap (Fold s -> a -> m s
fstep m s
initial s -> m b
extract) (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Maybe (s, s) -> m (Step (Maybe (s, s)) a)
step' forall a. Maybe a
Nothing

    where

    step' :: State Stream m a -> Maybe (s, s) -> m (Step (Maybe (s, s)) a)
step' State Stream m a
_ Maybe (s, s)
Nothing = do
        s
r <- m s
initial
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (s
r, s
state))

    step' State Stream m a
gst (Just (s
acc, s
st)) = s
acc seq :: forall a b. a -> b -> b
`seq` do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                s
acc' <- s -> a -> m s
fstep s
acc a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just (s
acc', s
s))
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (s
acc, s
s))
            Step s a
Stop      -> do
                forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ s -> m b
extract s
acc
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

{-# INLINE_NORMAL tapOffsetEvery #-}
tapOffsetEvery :: Monad m
    => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
tapOffsetEvery :: forall (m :: * -> *) a b.
Monad m =>
Int -> Int -> Fold m a b -> Stream m a -> Stream m a
tapOffsetEvery Int
offset Int
n (Fold s -> a -> m s
fstep m s
initial s -> m b
extract) (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> Maybe (s, s, Int) -> m (Step (Maybe (s, s, Int)) a)
step' forall a. Maybe a
Nothing

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> Maybe (s, s, Int) -> m (Step (Maybe (s, s, Int)) a)
step' State Stream m a
_ Maybe (s, s, Int)
Nothing = do
        s
r <- m s
initial
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (s
r, s
state, Int
offset forall a. Integral a => a -> a -> a
`mod` Int
n))

    step' State Stream m a
gst (Just (s
acc, s
st, Int
count)) | Int
count forall a. Ord a => a -> a -> Bool
<= Int
0 = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                !s
acc' <- s -> a -> m s
fstep s
acc a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just (s
acc', s
s, Int
n forall a. Num a => a -> a -> a
- Int
1))
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (s
acc, s
s, Int
count))
            Step s a
Stop      -> do
                forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ s -> m b
extract s
acc
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

    step' State Stream m a
gst (Just (s
acc, s
st, Int
count)) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just (s
acc, s
s, Int
count forall a. Num a => a -> a -> a
- Int
1))
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (s
acc, s
s, Int
count))
            Step s a
Stop      -> do
                forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ s -> m b
extract s
acc
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

{-# INLINE_NORMAL pollCounts #-}
pollCounts
    :: MonadAsync m
    => (a -> Bool)
    -> (Stream m Int -> Stream m Int)
    -> Fold m Int b
    -> Stream m a
    -> Stream m a
pollCounts :: forall (m :: * -> *) a b.
MonadAsync m =>
(a -> Bool)
-> (Stream m Int -> Stream m Int)
-> Fold m Int b
-> Stream m a
-> Stream m a
pollCounts a -> Bool
predicate Stream m Int -> Stream m Int
transf Fold m Int b
fld (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> Maybe (Var IO Int, ThreadId, s)
-> m (Step (Maybe (Var IO Int, ThreadId, s)) a)
step' forall a. Maybe a
Nothing
  where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> Maybe (Var IO Int, ThreadId, s)
-> m (Step (Maybe (Var IO Int, ThreadId, s)) a)
step' State Stream m a
_ Maybe (Var IO Int, ThreadId, s)
Nothing = do
        -- As long as we are using an "Int" for counts lockfree reads from
        -- Var should work correctly on both 32-bit and 64-bit machines.
        -- However, an Int on a 32-bit machine may overflow quickly.
        Var IO Int
countVar <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadMut m, Prim a) => a -> m (Var m a)
newVar (Int
0 :: Int)
        ThreadId
tid <- forall (m :: * -> *).
(MonadIO m, MonadBaseControl IO m) =>
m () -> m ThreadId
forkManaged
            forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
runFold Fold m Int b
fld
            forall a b. (a -> b) -> a -> b
$ Stream m Int -> Stream m Int
transf forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadIO m, Prim a) =>
Var IO a -> Stream m a
fromPrimVar Var IO Int
countVar
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (Var IO Int
countVar, ThreadId
tid, s
state))

    step' State Stream m a
gst (Just (Var IO Int
countVar, ThreadId
tid, s
st)) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (a -> Bool
predicate a
x) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadMut m, Prim a) =>
Var m a -> (a -> a) -> m ()
modifyVar' Var IO Int
countVar (forall a. Num a => a -> a -> a
+ Int
1)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just (Var IO Int
countVar, ThreadId
tid, s
s))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (Var IO Int
countVar, ThreadId
tid, s
s))
            Step s a
Stop -> do
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ThreadId -> IO ()
killThread ThreadId
tid
                forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL tapRate #-}
tapRate ::
       (MonadAsync m, MonadCatch m)
    => Double
    -> (Int -> m b)
    -> Stream m a
    -> Stream m a
tapRate :: forall (m :: * -> *) b a.
(MonadAsync m, MonadCatch m) =>
Double -> (Int -> m b) -> Stream m a -> Stream m a
tapRate Double
samplingRate Int -> m b
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> Maybe (Var IO Int, ThreadId, s, IORef ())
-> m (Step (Maybe (Var IO Int, ThreadId, s, IORef ())) a)
step' forall a. Maybe a
Nothing
  where
    {-# NOINLINE loop #-}
    loop :: Var IO Int -> Int -> m b
loop Var IO Int
countVar Int
prev = do
        Int
i <-
            forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
MC.catch
                (do forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Int -> IO ()
threadDelay (forall a b. (RealFrac a, Integral b) => a -> b
round forall a b. (a -> b) -> a -> b
$ Double
samplingRate forall a. Num a => a -> a -> a
* Double
1000000)
                    Int
i <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadMut m, Prim a) => Var m a -> m a
readVar Var IO Int
countVar
                    let !diff :: Int
diff = Int
i forall a. Num a => a -> a -> a
- Int
prev
                    forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ Int -> m b
action Int
diff
                    forall (m :: * -> *) a. Monad m => a -> m a
return Int
i)
                (\(AsyncException
e :: AsyncException) -> do
                     Int
i <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadMut m, Prim a) => Var m a -> m a
readVar Var IO Int
countVar
                     let !diff :: Int
diff = Int
i forall a. Num a => a -> a -> a
- Int
prev
                     forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ Int -> m b
action Int
diff
                     forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (forall e. Exception e => e -> SomeException
MC.toException AsyncException
e))
        Var IO Int -> Int -> m b
loop Var IO Int
countVar Int
i

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> Maybe (Var IO Int, ThreadId, s, IORef ())
-> m (Step (Maybe (Var IO Int, ThreadId, s, IORef ())) a)
step' State Stream m a
_ Maybe (Var IO Int, ThreadId, s, IORef ())
Nothing = do
        Var IO Int
countVar <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadMut m, Prim a) => a -> m (Var m a)
newVar Int
0
        ThreadId
tid <- forall (m :: * -> *). MonadBaseControl IO m => m () -> m ThreadId
fork forall a b. (a -> b) -> a -> b
$ forall {b}. Var IO Int -> Int -> m b
loop Var IO Int
countVar Int
0
        IORef ()
ref <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. a -> IO (IORef a)
newIORef ()
        Weak (IORef ())
_ <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef ()
ref (ThreadId -> IO ()
killThread ThreadId
tid)
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (Var IO Int
countVar, ThreadId
tid, s
state, IORef ()
ref))

    step' State Stream m a
gst (Just (Var IO Int
countVar, ThreadId
tid, s
st, IORef ()
ref)) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a.
(MonadMut m, Prim a) =>
Var m a -> (a -> a) -> m ()
modifyVar' Var IO Int
countVar (forall a. Num a => a -> a -> a
+ Int
1)
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just (Var IO Int
countVar, ThreadId
tid, s
s, IORef ()
ref))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just (Var IO Int
countVar, ThreadId
tid, s
s, IORef ()
ref))
            Step s a
Stop -> do
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ThreadId -> IO ()
killThread ThreadId
tid
                forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop


-------------------------------------------------------------------------------
-- Filtering
-------------------------------------------------------------------------------

{-# INLINE_NORMAL takeWhileM #-}
takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
takeWhileM :: forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
takeWhileM a -> m Bool
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step' s
state
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Bool
b <- a -> m Bool
f a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if Bool
b then forall s a. a -> s -> Step s a
Yield a
x s
s else forall s a. Step s a
Stop
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE takeWhile #-}
takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
takeWhile :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile a -> Bool
f = forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
takeWhileM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

{-# INLINE_NORMAL drop #-}
drop :: Monad m => Int -> Stream m a -> Stream m a
drop :: forall (m :: * -> *) a. Monad m => Int -> Stream m a -> Stream m a
drop Int
n (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {a}.
(Ord a, Num a) =>
State Stream m a -> (s, Maybe a) -> m (Step (s, Maybe a) a)
step' (s
state, forall a. a -> Maybe a
Just Int
n)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, Maybe a) -> m (Step (s, Maybe a) a)
step' State Stream m a
gst (s
st, Just a
i)
      | a
i forall a. Ord a => a -> a -> Bool
> a
0 = do
          Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
          forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            case Step s a
r of
              Yield a
_ s
s -> forall s a. s -> Step s a
Skip (s
s, forall a. a -> Maybe a
Just (a
i forall a. Num a => a -> a -> a
- a
1))
              Skip s
s    -> forall s a. s -> Step s a
Skip (s
s, forall a. a -> Maybe a
Just a
i)
              Step s a
Stop      -> forall s a. Step s a
Stop
      | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
st, forall a. Maybe a
Nothing)

    step' State Stream m a
gst (s
st, Maybe a
Nothing) = do
      Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
        case Step s a
r of
          Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x (s
s, forall a. Maybe a
Nothing)
          Skip  s
s   -> forall s a. s -> Step s a
Skip (s
s, forall a. Maybe a
Nothing)
          Step s a
Stop      -> forall s a. Step s a
Stop

data DropWhileState s a
    = DropWhileDrop s
    | DropWhileYield a s
    | DropWhileNext s

{-# INLINE_NORMAL dropWhileM #-}
dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
dropWhileM :: forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
dropWhileM a -> m Bool
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {a}.
State Stream m a
-> DropWhileState s a -> m (Step (DropWhileState s a) a)
step' (forall s a. s -> DropWhileState s a
DropWhileDrop s
state)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> DropWhileState s a -> m (Step (DropWhileState s a) a)
step' State Stream m a
gst (DropWhileDrop s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Bool
b <- a -> m Bool
f a
x
                if Bool
b
                then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> DropWhileState s a
DropWhileDrop s
s)
                else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. a -> s -> DropWhileState s a
DropWhileYield a
x s
s)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> DropWhileState s a
DropWhileDrop s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
gst (DropWhileNext s
st) =  do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. a -> s -> DropWhileState s a
DropWhileYield a
x s
s)
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> DropWhileState s a
DropWhileNext s
s)
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
_ (DropWhileYield a
x s
st) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall s a. s -> DropWhileState s a
DropWhileNext s
st)

{-# INLINE dropWhile #-}
dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
dropWhile :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
dropWhile a -> Bool
f = forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
dropWhileM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

{-# INLINE_NORMAL filterM #-}
filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
filterM :: forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
filterM a -> m Bool
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step' s
state
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Bool
b <- a -> m Bool
f a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if Bool
b
                         then forall s a. a -> s -> Step s a
Yield a
x s
s
                         else forall s a. s -> Step s a
Skip s
s
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE filter #-}
filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
filter :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
filter a -> Bool
f = forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
filterM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

{-# INLINE_NORMAL uniq #-}
uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
uniq :: forall a (m :: * -> *). (Eq a, Monad m) => Stream m a -> Stream m a
uniq (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> (Maybe a, s) -> m (Step (Maybe a, s) a)
step' (forall a. Maybe a
Nothing, s
state)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (Maybe a, s) -> m (Step (Maybe a, s) a)
step' State Stream m a
gst (Maybe a
Nothing, s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just a
x, s
s)
            Skip  s
s   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip  (forall a. Maybe a
Nothing, s
s)
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step' State Stream m a
gst (Just a
x, s
st)  = do
         Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
         case Step s a
r of
             Yield a
y s
s | a
x forall a. Eq a => a -> a -> Bool
== a
y   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just a
x, s
s)
                       | Bool
otherwise -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
y (forall a. a -> Maybe a
Just a
y, s
s)
             Skip  s
s   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just a
x, s
s)
             Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Transformation by Mapping
------------------------------------------------------------------------------

{-# INLINE_NORMAL sequence #-}
sequence :: Monad m => Stream m (m a) -> Stream m a
sequence :: forall (m :: * -> *) a. Monad m => Stream m (m a) -> Stream m a
sequence (Stream State Stream m (m a) -> s -> m (Step s (m a))
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}. State Stream m a -> s -> m (Step s a)
step' s
state
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst s
st = do
         Step s (m a)
r <- State Stream m (m a) -> s -> m (Step s (m a))
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
         case Step s (m a)
r of
             Yield m a
x s
s -> m a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
a -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. a -> s -> Step s a
Yield a
a s
s)
             Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip s
s
             Step s (m a)
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Inserting
------------------------------------------------------------------------------

data LoopState x s = FirstYield s
                   | InterspersingYield s
                   | YieldAndCarry x s

{-# INLINE_NORMAL intersperseM #-}
intersperseM :: Monad m => m a -> Stream m a -> Stream m a
intersperseM :: forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
intersperseM m a
m (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> LoopState a s -> m (Step (LoopState a s) a)
step' (forall x s. s -> LoopState x s
FirstYield s
state)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> LoopState a s -> m (Step (LoopState a s) a)
step' State Stream m a
gst (FirstYield s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            case Step s a
r of
                Yield a
x s
s -> forall s a. s -> Step s a
Skip (forall x s. x -> s -> LoopState x s
YieldAndCarry a
x s
s)
                Skip s
s -> forall s a. s -> Step s a
Skip (forall x s. s -> LoopState x s
FirstYield s
s)
                Step s a
Stop -> forall s a. Step s a
Stop

    step' State Stream m a
gst (InterspersingYield s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                a
a <- m a
m
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (forall x s. x -> s -> LoopState x s
YieldAndCarry a
x s
s)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall x s. s -> LoopState x s
InterspersingYield s
s
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
_ (YieldAndCarry a
x s
st) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (forall x s. s -> LoopState x s
InterspersingYield s
st)

data SuffixState s a
    = SuffixElem s
    | SuffixSuffix s
    | SuffixYield a (SuffixState s a)

{-# INLINE_NORMAL intersperseSuffix #-}
intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
intersperseSuffix :: forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
intersperseSuffix m a
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> SuffixState s a -> m (Step (SuffixState s a) a)
step' (forall s a. s -> SuffixState s a
SuffixElem s
state)
    where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> SuffixState s a -> m (Step (SuffixState s a) a)
step' State Stream m a
gst (SuffixElem s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. s -> Step s a
Skip (forall s a. a -> SuffixState s a -> SuffixState s a
SuffixYield a
x (forall s a. s -> SuffixState s a
SuffixSuffix s
s))
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s a. s -> SuffixState s a
SuffixElem s
s)
            Step s a
Stop -> forall s a. Step s a
Stop

    step' State Stream m a
_ (SuffixSuffix s
st) = do
        m a
action forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. a -> SuffixState s a -> SuffixState s a
SuffixYield a
r (forall s a. s -> SuffixState s a
SuffixElem s
st))

    step' State Stream m a
_ (SuffixYield a
x SuffixState s a
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x SuffixState s a
next

data SuffixSpanState s a
    = SuffixSpanElem s Int
    | SuffixSpanSuffix s
    | SuffixSpanYield a (SuffixSpanState s a)
    | SuffixSpanLast
    | SuffixSpanStop

-- | intersperse after every n items
{-# INLINE_NORMAL intersperseSuffixBySpan #-}
intersperseSuffixBySpan :: forall m a. Monad m
    => Int -> m a -> Stream m a -> Stream m a
intersperseSuffixBySpan :: forall (m :: * -> *) a.
Monad m =>
Int -> m a -> Stream m a -> Stream m a
intersperseSuffixBySpan Int
n m a
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> SuffixSpanState s a -> m (Step (SuffixSpanState s a) a)
step' (forall s a. s -> Int -> SuffixSpanState s a
SuffixSpanElem s
state Int
n)
    where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> SuffixSpanState s a -> m (Step (SuffixSpanState s a) a)
step' State Stream m a
gst (SuffixSpanElem s
st Int
i) | Int
i forall a. Ord a => a -> a -> Bool
> Int
0 = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
x s
s -> forall s a. s -> Step s a
Skip (forall s a. a -> SuffixSpanState s a -> SuffixSpanState s a
SuffixSpanYield a
x (forall s a. s -> Int -> SuffixSpanState s a
SuffixSpanElem s
s (Int
i forall a. Num a => a -> a -> a
- Int
1)))
            Skip s
s -> forall s a. s -> Step s a
Skip (forall s a. s -> Int -> SuffixSpanState s a
SuffixSpanElem s
s Int
i)
            Step s a
Stop -> if Int
i forall a. Eq a => a -> a -> Bool
== Int
n then forall s a. Step s a
Stop else forall s a. s -> Step s a
Skip forall s a. SuffixSpanState s a
SuffixSpanLast
    step' State Stream m a
_ (SuffixSpanElem s
st Int
_) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. s -> SuffixSpanState s a
SuffixSpanSuffix s
st)

    step' State Stream m a
_ (SuffixSpanSuffix s
st) = do
        m a
action forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. a -> SuffixSpanState s a -> SuffixSpanState s a
SuffixSpanYield a
r (forall s a. s -> Int -> SuffixSpanState s a
SuffixSpanElem s
st Int
n))

    step' State Stream m a
_ (SuffixSpanState s a
SuffixSpanLast) = do
        m a
action forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \a
r -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s a. a -> SuffixSpanState s a -> SuffixSpanState s a
SuffixSpanYield a
r forall s a. SuffixSpanState s a
SuffixSpanStop)

    step' State Stream m a
_ (SuffixSpanYield a
x SuffixSpanState s a
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x SuffixSpanState s a
next

    step' State Stream m a
_ (SuffixSpanState s a
SuffixSpanStop) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE intersperse #-}
intersperse :: Monad m => a -> Stream m a -> Stream m a
intersperse :: forall (m :: * -> *) a. Monad m => a -> Stream m a -> Stream m a
intersperse a
a = forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
intersperseM (forall (m :: * -> *) a. Monad m => a -> m a
return a
a)

{-# INLINE_NORMAL insertBy #-}
insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
insertBy :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> Ordering) -> a -> Stream m a -> Stream m a
insertBy a -> a -> Ordering
cmp a
a (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> (s, Bool, Maybe a) -> m (Step (s, Bool, Maybe a) a)
step' (s
state, Bool
False, forall a. Maybe a
Nothing)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> (s, Bool, Maybe a) -> m (Step (s, Bool, Maybe a) a)
step' State Stream m a
gst (s
st, Bool
False, Maybe a
_) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> case a -> a -> Ordering
cmp a
a a
x of
                Ordering
GT -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x (s
s, Bool
False, forall a. Maybe a
Nothing)
                Ordering
_  -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (s
s, Bool
True, forall a. a -> Maybe a
Just a
x)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, Bool
False, forall a. Maybe a
Nothing)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (s
st, Bool
True, forall a. Maybe a
Nothing)

    step' State Stream m a
_ (s
_, Bool
True, Maybe a
Nothing) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
gst (s
st, Bool
True, Just a
prev) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
prev (s
s, Bool
True, forall a. a -> Maybe a
Just a
x)
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, Bool
True, forall a. a -> Maybe a
Just a
prev)
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
prev (s
st, Bool
True, forall a. Maybe a
Nothing)

------------------------------------------------------------------------------
-- Deleting
------------------------------------------------------------------------------

{-# INLINE_NORMAL deleteBy #-}
deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
deleteBy :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> Bool) -> a -> Stream m a -> Stream m a
deleteBy a -> a -> Bool
eq a
x (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> (s, Bool) -> m (Step (s, Bool) a)
step' (s
state, Bool
False)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, Bool) -> m (Step (s, Bool) a)
step' State Stream m a
gst (s
st, Bool
False) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
y s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
                if a -> a -> Bool
eq a
x a
y then forall s a. s -> Step s a
Skip (s
s, Bool
True) else forall s a. a -> s -> Step s a
Yield a
y (s
s, Bool
False)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, Bool
False)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
gst (s
st, Bool
True) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
y s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
y (s
s, Bool
True)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, Bool
True)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Transformation by Map and Filter
------------------------------------------------------------------------------

-- XXX Will this always fuse properly?
{-# INLINE_NORMAL mapMaybe #-}
mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
mapMaybe :: forall (m :: * -> *) a b.
Monad m =>
(a -> Maybe b) -> Stream m a -> Stream m b
mapMaybe a -> Maybe b
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. HasCallStack => Maybe a -> a
fromJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
filter forall a. Maybe a -> Bool
isJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
map a -> Maybe b
f

{-# INLINE_NORMAL mapMaybeM #-}
mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
mapMaybeM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m (Maybe b)) -> Stream m a -> Stream m b
mapMaybeM a -> m (Maybe b)
f = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. HasCallStack => Maybe a -> a
fromJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
filter forall a. Maybe a -> Bool
isJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
mapM a -> m (Maybe b)
f

------------------------------------------------------------------------------
-- Zipping
------------------------------------------------------------------------------

{-# INLINE_NORMAL indexed #-}
indexed :: Monad m => Stream m a -> Stream m (Int, a)
indexed :: forall (m :: * -> *) a. Monad m => Stream m a -> Stream m (Int, a)
indexed (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {b} {m :: * -> *} {a}.
Num b =>
State Stream m a -> (s, b) -> m (Step (s, b) (b, a))
step' (s
state, Int
0)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, b) -> m (Step (s, b) (b, a))
step' State Stream m a
gst (s
st, b
i) = b
i seq :: forall a b. a -> b -> b
`seq` do
         Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
         case Step s a
r of
             Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield (b
i, a
x) (s
s, b
iforall a. Num a => a -> a -> a
+b
1)
             Skip    s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, b
i)
             Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL indexedR #-}
indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
indexedR :: forall (m :: * -> *) a.
Monad m =>
Int -> Stream m a -> Stream m (Int, a)
indexedR Int
m (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {b} {m :: * -> *} {a}.
Num b =>
State Stream m a -> (s, b) -> m (Step (s, b) (b, a))
step' (s
state, Int
m)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, b) -> m (Step (s, b) (b, a))
step' State Stream m a
gst (s
st, b
i) = b
i seq :: forall a b. a -> b -> b
`seq` do
         Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
         case Step s a
r of
             Yield a
x s
s -> let i' :: b
i' = b
i forall a. Num a => a -> a -> a
- b
1
                          in forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield (b
i, a
x) (s
s, b
i')
             Skip    s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
s, b
i)
             Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE_NORMAL zipWithM #-}
zipWithM :: Monad m
    => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
zipWithM :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
zipWithM a -> b -> m c
f (Stream State Stream m a -> s -> m (Step s a)
stepa s
ta) (Stream State Stream m b -> s -> m (Step s b)
stepb s
tb) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a -> (s, s, Maybe a) -> m (Step (s, s, Maybe a) c)
step (s
ta, s
tb, forall a. Maybe a
Nothing)
  where
    {-# INLINE_LATE step #-}
    step :: State Stream m a -> (s, s, Maybe a) -> m (Step (s, s, Maybe a) c)
step State Stream m a
gst (s
sa, s
sb, Maybe a
Nothing) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepa (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
sa
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
          case Step s a
r of
            Yield a
x s
sa' -> forall s a. s -> Step s a
Skip (s
sa', s
sb, forall a. a -> Maybe a
Just a
x)
            Skip s
sa'    -> forall s a. s -> Step s a
Skip (s
sa', s
sb, forall a. Maybe a
Nothing)
            Step s a
Stop        -> forall s a. Step s a
Stop

    step State Stream m a
gst (s
sa, s
sb, Just a
x) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
stepb (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
sb
        case Step s b
r of
            Yield b
y s
sb' -> do
                c
z <- a -> b -> m c
f a
x b
y
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield c
z (s
sa, s
sb', forall a. Maybe a
Nothing)
            Skip s
sb' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (s
sa, s
sb', forall a. a -> Maybe a
Just a
x)
            Step s b
Stop     -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

#if __GLASGOW_HASKELL__ >= 801
{-# RULES "zipWithM xs xs"
    forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
#endif

{-# INLINE zipWith #-}
zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
zipWith :: forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
zipWith a -> b -> c
f = forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
zipWithM (\a
a b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return (a -> b -> c
f a
a b
b))

------------------------------------------------------------------------------
-- Merging
------------------------------------------------------------------------------

{-# INLINE_NORMAL mergeByM #-}
mergeByM
    :: (Monad m)
    => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
mergeByM :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
mergeByM a -> a -> m Ordering
cmp (Stream State Stream m a -> s -> m (Step s a)
stepa s
ta) (Stream State Stream m a -> s -> m (Step s a)
stepb s
tb) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> (Maybe s, Maybe s, Maybe a, Maybe a)
-> m (Step (Maybe s, Maybe s, Maybe a, Maybe a) a)
step (forall a. a -> Maybe a
Just s
ta, forall a. a -> Maybe a
Just s
tb, forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)
  where
    {-# INLINE_LATE step #-}

    -- one of the values is missing, and the corresponding stream is running
    step :: State Stream m a
-> (Maybe s, Maybe s, Maybe a, Maybe a)
-> m (Step (Maybe s, Maybe s, Maybe a, Maybe a) a)
step State Stream m a
gst (Just s
sa, Maybe s
sb, Maybe a
Nothing, Maybe a
b) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepa State Stream m a
gst s
sa
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
sa' -> forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
sa', Maybe s
sb, forall a. a -> Maybe a
Just a
a, Maybe a
b)
            Skip s
sa'    -> forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
sa', Maybe s
sb, forall a. Maybe a
Nothing, Maybe a
b)
            Step s a
Stop        -> forall s a. s -> Step s a
Skip (forall a. Maybe a
Nothing, Maybe s
sb, forall a. Maybe a
Nothing, Maybe a
b)

    step State Stream m a
gst (Maybe s
sa, Just s
sb, Maybe a
a, Maybe a
Nothing) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
stepb State Stream m a
gst s
sb
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
b s
sb' -> forall s a. s -> Step s a
Skip (Maybe s
sa, forall a. a -> Maybe a
Just s
sb', Maybe a
a, forall a. a -> Maybe a
Just a
b)
            Skip s
sb'    -> forall s a. s -> Step s a
Skip (Maybe s
sa, forall a. a -> Maybe a
Just s
sb', Maybe a
a, forall a. Maybe a
Nothing)
            Step s a
Stop        -> forall s a. s -> Step s a
Skip (Maybe s
sa, forall a. Maybe a
Nothing, Maybe a
a, forall a. Maybe a
Nothing)

    -- both the values are available
    step State Stream m a
_ (Maybe s
sa, Maybe s
sb, Just a
a, Just a
b) = do
        Ordering
res <- a -> a -> m Ordering
cmp a
a a
b
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Ordering
res of
            Ordering
GT -> forall s a. a -> s -> Step s a
Yield a
b (Maybe s
sa, Maybe s
sb, forall a. a -> Maybe a
Just a
a, forall a. Maybe a
Nothing)
            Ordering
_  -> forall s a. a -> s -> Step s a
Yield a
a (Maybe s
sa, Maybe s
sb, forall a. Maybe a
Nothing, forall a. a -> Maybe a
Just a
b)

    -- one of the values is missing, corresponding stream is done
    step State Stream m a
_ (Maybe s
Nothing, Maybe s
sb, Maybe a
Nothing, Just a
b) =
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
b (forall a. Maybe a
Nothing, Maybe s
sb, forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

    step State Stream m a
_ (Maybe s
sa, Maybe s
Nothing, Just a
a, Maybe a
Nothing) =
            forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (Maybe s
sa, forall a. Maybe a
Nothing, forall a. Maybe a
Nothing, forall a. Maybe a
Nothing)

    step State Stream m a
_ (Maybe s
Nothing, Maybe s
Nothing, Maybe a
Nothing, Maybe a
Nothing) = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE mergeBy #-}
mergeBy
    :: (Monad m)
    => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
mergeBy :: forall (m :: * -> *) a.
Monad m =>
(a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
mergeBy a -> a -> Ordering
cmp = forall (m :: * -> *) a.
Monad m =>
(a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
mergeByM (\a
a a
b -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ a -> a -> Ordering
cmp a
a a
b)

------------------------------------------------------------------------------
-- Transformation comprehensions
------------------------------------------------------------------------------

{-# INLINE_NORMAL the #-}
the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
the :: forall a (m :: * -> *).
(Eq a, Monad m) =>
Stream m a -> m (Maybe a)
the (Stream State Stream m a -> s -> m (Step s a)
step s
state) = s -> m (Maybe a)
go s
state
  where
    go :: s -> m (Maybe a)
go s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> a -> s -> m (Maybe a)
go' a
x s
s
            Skip s
s    -> s -> m (Maybe a)
go s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
    go' :: a -> s -> m (Maybe a)
go' a
n s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s | a
x forall a. Eq a => a -> a -> Bool
== a
n -> a -> s -> m (Maybe a)
go' a
n s
s
                      | Bool
otherwise -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
            Skip s
s -> a -> s -> m (Maybe a)
go' a
n s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return (forall a. a -> Maybe a
Just a
n)

{-# INLINE runFold #-}
runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
runFold :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
runFold (Fold s -> a -> m s
step m s
begin s -> m b
done) = forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
foldlMx' s -> a -> m s
step m s
begin s -> m b
done

-------------------------------------------------------------------------------
-- Concurrent application and fold
-------------------------------------------------------------------------------

-- XXX These functions should be moved to Stream/Parallel.hs
--
-- Using StreamD the worker stream producing code can fuse with the code to
-- queue output to the SVar giving some perf boost.
--
-- Note that StreamD can only be used in limited situations, specifically, we
-- cannot implement joinStreamVarPar using this.
--
-- XXX make sure that the SVar passed is a Parallel style SVar.

-- | Fold the supplied stream to the SVar asynchronously using Parallel
-- concurrency style.
-- {-# INLINE_NORMAL toSVarParallel #-}
{-# INLINE toSVarParallel #-}
toSVarParallel :: MonadAsync m
    => State t m a -> SVar t m a -> Stream m a -> m ()
toSVarParallel :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
State t m a -> SVar t m a -> Stream m a -> m ()
toSVarParallel State t m a
st SVar t m a
sv Stream m a
xs =
    if forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> Bool
svarInspectMode SVar t m a
sv
    then m ()
forkWithDiag
    else do
        ThreadId
tid <-
                case forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
State t m a -> Maybe Count
getYieldLimit State t m a
st of
                    Maybe Count
Nothing -> forall (m :: * -> *).
MonadBaseControl IO m =>
m () -> RunInIO m -> (SomeException -> IO ()) -> m ThreadId
doFork (Maybe WorkerInfo -> m ()
work forall a. Maybe a
Nothing)
                                      (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> RunInIO m
svarMrun SVar t m a
sv)
                                      (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SomeException -> IO ()
handleChildException SVar t m a
sv)
                    Just Count
_  -> forall (m :: * -> *).
MonadBaseControl IO m =>
m () -> RunInIO m -> (SomeException -> IO ()) -> m ThreadId
doFork (Maybe WorkerInfo -> m ()
workLim forall a. Maybe a
Nothing)
                                      (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> RunInIO m
svarMrun SVar t m a
sv)
                                      (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SomeException -> IO ()
handleChildException SVar t m a
sv)
        forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> ThreadId -> m ()
modifyThread SVar t m a
sv ThreadId
tid

    where

    {-# NOINLINE work #-}
    work :: Maybe WorkerInfo -> m ()
work Maybe WorkerInfo
info = (forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
runFold (forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> Maybe WorkerInfo -> Fold m a ()
FL.toParallelSVar SVar t m a
sv Maybe WorkerInfo
info) Stream m a
xs)

    {-# NOINLINE workLim #-}
    workLim :: Maybe WorkerInfo -> m ()
workLim Maybe WorkerInfo
info = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
runFold (forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> Maybe WorkerInfo -> Fold m a ()
FL.toParallelSVarLimited SVar t m a
sv Maybe WorkerInfo
info) Stream m a
xs

    {-# NOINLINE forkWithDiag #-}
    forkWithDiag :: m ()
forkWithDiag = do
        -- We do not use workerCount in case of ParallelVar but still there is
        -- no harm in maintaining it correctly.
        forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall t. IORef t -> (t -> t) -> IO ()
atomicModifyIORefCAS_ (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> IORef Int
workerCount SVar t m a
sv) forall a b. (a -> b) -> a -> b
$ \Int
n -> Int
n forall a. Num a => a -> a -> a
+ Int
1
        forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> m ()
recordMaxWorkers SVar t m a
sv
        -- This allocation matters when significant number of workers are being
        -- sent. We allocate it only when needed. The overhead increases by 4x.
        Maybe WorkerInfo
winfo <-
            case forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> Maybe YieldRateInfo
yieldRateInfo SVar t m a
sv of
                Maybe YieldRateInfo
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing
                Just YieldRateInfo
_ -> forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ do
                    IORef Count
cntRef <- forall a. a -> IO (IORef a)
newIORef Count
0
                    AbsTime
t <- Clock -> IO AbsTime
getTime Clock
Monotonic
                    IORef (Count, AbsTime)
lat <- forall a. a -> IO (IORef a)
newIORef (Count
0, AbsTime
t)
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just WorkerInfo
                        { workerYieldMax :: Count
workerYieldMax = Count
0
                        , workerYieldCount :: IORef Count
workerYieldCount = IORef Count
cntRef
                        , workerLatencyStart :: IORef (Count, AbsTime)
workerLatencyStart = IORef (Count, AbsTime)
lat
                        }
        ThreadId
tid <-
            case forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
State t m a -> Maybe Count
getYieldLimit State t m a
st of
                Maybe Count
Nothing -> forall (m :: * -> *).
MonadBaseControl IO m =>
m () -> RunInIO m -> (SomeException -> IO ()) -> m ThreadId
doFork (Maybe WorkerInfo -> m ()
work Maybe WorkerInfo
winfo)
                                  (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> RunInIO m
svarMrun SVar t m a
sv)
                                  (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SomeException -> IO ()
handleChildException SVar t m a
sv)
                Just Count
_  -> forall (m :: * -> *).
MonadBaseControl IO m =>
m () -> RunInIO m -> (SomeException -> IO ()) -> m ThreadId
doFork (Maybe WorkerInfo -> m ()
workLim Maybe WorkerInfo
winfo)
                                  (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> RunInIO m
svarMrun SVar t m a
sv)
                                  (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SomeException -> IO ()
handleChildException SVar t m a
sv)
        forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> ThreadId -> m ()
modifyThread SVar t m a
sv ThreadId
tid

{-# INLINE_NORMAL mkParallelD #-}
mkParallelD :: MonadAsync m => Stream m a -> Stream m a
mkParallelD :: forall (m :: * -> *) a. MonadAsync m => Stream m a -> Stream m a
mkParallelD Stream m a
m = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> Maybe (Stream m a) -> m (Step (Maybe (Stream m a)) a)
step forall a. Maybe a
Nothing
    where

    step :: State Stream m a
-> Maybe (Stream m a) -> m (Step (Maybe (Stream m a)) a)
step State Stream m a
gst Maybe (Stream m a)
Nothing = do
        SVar Stream m a
sv <- forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVarStopStyle -> State t m a -> m (SVar t m a)
newParallelVar SVarStopStyle
StopNone State Stream m a
gst
        forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
State t m a -> SVar t m a -> Stream m a -> m ()
toSVarParallel State Stream m a
gst SVar Stream m a
sv Stream m a
m
        -- XXX use unfold instead?
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVar t m a -> Stream m a
fromSVar SVar Stream m a
sv

    step State Stream m a
gst (Just (UnStream State Stream m a -> s -> m (Step s a)
step1 s
st)) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step1 s
s)
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step1 s
s)
            Step s a
Stop      -> forall s a. Step s a
Stop

-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
-- parallel style SVar for evaluation. Currently, parallel style cannot use
-- rate control whereas Async style can use rate control. In async style SVar
-- the worker thread terminates when the buffer is full whereas in Parallel
-- style it blocks.
--
-- | Make the stream producer and consumer run concurrently by introducing a
-- buffer between them. The producer thread evaluates the input stream until
-- the buffer fills, it blocks if the buffer is full until there is space in
-- the buffer. The consumer consumes the stream lazily from the buffer.
--
-- /Internal/
--
{-# INLINE_NORMAL mkParallel #-}
mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
mkParallel :: forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, MonadAsync m) =>
t m a -> t m a
mkParallel = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
fromStreamD forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. MonadAsync m => Stream m a -> Stream m a
mkParallelD forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
t m a -> Stream m a
toStreamD

-- Note: we can use another API with two callbacks stop and yield if we want
-- the callback to be able to indicate end of stream.
--
-- | Generates a callback and a stream pair. The callback returned is used to
-- queue values to the stream.  The stream is infinite, there is no way for the
-- callback to indicate that it is done now.
--
-- /Internal/
--
{-# INLINE_NORMAL newCallbackStream #-}
newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
newCallbackStream :: forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, MonadAsync m) =>
m (a -> m (), t m a)
newCallbackStream = do
    SVar Any m a
sv <- forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVarStopStyle -> State t m a -> m (SVar t m a)
newParallelVar SVarStopStyle
StopNone forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState

    -- XXX Add our own thread-id to the SVar as we can not know the callback's
    -- thread-id and the callback is not run in a managed worker. We need to
    -- handle this better.
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ThreadId
myThreadId forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> ThreadId -> m ()
modifyThread SVar Any m a
sv

    let callback :: a -> m ()
callback a
a = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> ChildEvent a -> IO Int
send SVar Any m a
sv (forall a. a -> ChildEvent a
ChildYield a
a)
    -- XXX we can return an SVar and then the consumer can unfold from the
    -- SVar?
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall {m :: * -> *}. MonadIO m => a -> m ()
callback, forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(IsStream t, Monad m) =>
Stream m a -> t m a
fromStreamD (forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVar t m a -> Stream m a
fromSVar SVar Any m a
sv))

-------------------------------------------------------------------------------
-- Concurrent tap
-------------------------------------------------------------------------------

-- | Create an SVar with a fold consumer that will fold any elements sent to it
-- using the supplied fold function.
{-# INLINE newFoldSVar #-}
newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
newFoldSVar :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
MonadAsync m =>
State t m a -> Fold m a b -> m (SVar t m a)
newFoldSVar State t m a
stt Fold m a b
f = do
    -- Buffer size for the SVar is derived from the current state
    SVar t m a
sv <- forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVarStopStyle -> State t m a -> m (SVar t m a)
newParallelVar SVarStopStyle
StopAny (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State t m a
stt)
    -- Add the producer thread-id to the SVar.
    forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO ThreadId
myThreadId forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadIO m =>
SVar t m a -> ThreadId -> m ()
modifyThread SVar t m a
sv
    forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
MonadBaseControl IO m =>
m () -> RunInIO m -> (SomeException -> IO ()) -> m ThreadId
doFork (forall {t :: (* -> *) -> * -> *}. SVar t m a -> m ()
work SVar t m a
sv) (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> RunInIO m
svarMrun SVar t m a
sv) (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> SomeException -> IO ()
handleFoldException SVar t m a
sv)
    forall (m :: * -> *) a. Monad m => a -> m a
return SVar t m a
sv

    where

    {-# NOINLINE work #-}
    work :: SVar t m a -> m ()
work SVar t m a
sv = forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
runFold Fold m a b
f forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) (t :: (* -> *) -> * -> *) a.
MonadAsync m =>
SVar t m a -> Stream m a
fromProducer SVar t m a
sv

data TapState sv st = TapInit | Tapping sv st | TapDone st

{-# INLINE_NORMAL tapAsync #-}
tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
tapAsync :: forall (m :: * -> *) a b.
MonadAsync m =>
Fold m a b -> Stream m a -> Stream m a
tapAsync Fold m a b
f (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> TapState (SVar Stream m a) s
-> m (Step (TapState (SVar Stream m a) s) a)
step forall sv st. TapState sv st
TapInit
    where

    drainFold :: SVar Stream m a -> m ()
drainFold SVar Stream m a
svr = do
            -- In general, a Stop event would come equipped with the result
            -- of the fold. It is not used here but it would be useful in
            -- applicative and distribute.
            Bool
done <- forall (m :: * -> *) a. MonadAsync m => SVar Stream m a -> m Bool
fromConsumer SVar Stream m a
svr
            forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
done) forall a b. (a -> b) -> a -> b
$ do
                forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> String -> IO () -> IO ()
withDiagMVar SVar Stream m a
svr String
"teeToSVar: waiting to drain"
                       forall a b. (a -> b) -> a -> b
$ forall a. MVar a -> IO a
takeMVar (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> MVar ()
outputDoorBellFromConsumer SVar Stream m a
svr)
                SVar Stream m a -> m ()
drainFold SVar Stream m a
svr

    stopFold :: SVar Stream m a -> m ()
stopFold SVar Stream m a
svr = do
            forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
SVar t m a -> Maybe WorkerInfo -> IO ()
sendStop SVar Stream m a
svr forall a. Maybe a
Nothing
            -- drain/wait until a stop event arrives from the fold.
            forall {m :: * -> *} {a}.
(MonadIO m, MonadBaseControl IO m, MonadThrow m) =>
SVar Stream m a -> m ()
drainFold SVar Stream m a
svr

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> TapState (SVar Stream m a) s
-> m (Step (TapState (SVar Stream m a) s) a)
step State Stream m a
gst TapState (SVar Stream m a) s
TapInit = do
        SVar Stream m a
sv <- forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
MonadAsync m =>
State t m a -> Fold m a b -> m (SVar t m a)
newFoldSVar State Stream m a
gst Fold m a b
f
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall sv st. sv -> st -> TapState sv st
Tapping SVar Stream m a
sv s
state1)

    step State Stream m a
gst (Tapping SVar Stream m a
sv s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        case Step s a
r of
            Yield a
a s
s ->  do
                Bool
done <- forall (m :: * -> *) a.
MonadAsync m =>
SVar Stream m a -> a -> m Bool
pushToFold SVar Stream m a
sv a
a
                if Bool
done
                then do
                    -- XXX we do not need to wait synchronously here
                    forall {m :: * -> *} {a}.
(MonadIO m, MonadBaseControl IO m, MonadThrow m) =>
SVar Stream m a -> m ()
stopFold SVar Stream m a
sv
                    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (forall sv st. st -> TapState sv st
TapDone s
s)
                else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a (forall sv st. sv -> st -> TapState sv st
Tapping SVar Stream m a
sv s
s)
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall sv st. sv -> st -> TapState sv st
Tapping SVar Stream m a
sv s
s)
            Step s a
Stop -> do
                forall {m :: * -> *} {a}.
(MonadIO m, MonadBaseControl IO m, MonadThrow m) =>
SVar Stream m a -> m ()
stopFold SVar Stream m a
sv
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. Step s a
Stop

    step State Stream m a
gst (TapDone s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall sv st. st -> TapState sv st
TapDone s
s)
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall sv st. st -> TapState sv st
TapDone s
s)
            Step s a
Stop      -> forall s a. Step s a
Stop

-- XXX Exported from Array again as this fold is specific to Array
-- | Take last 'n' elements from the stream and discard the rest.
{-# INLINE lastN #-}
lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
lastN :: forall a (m :: * -> *).
(Storable a, MonadIO m) =>
Int -> Fold m a (Array a)
lastN Int
n
    | Int
n forall a. Ord a => a -> a -> Bool
<= Int
0 = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b. a -> b -> a
const forall a. Monoid a => a
mempty) forall (m :: * -> *) a. Monad m => Fold m a ()
FL.drain
    | Bool
otherwise = forall (m :: * -> *) a b s.
(s -> a -> m s) -> m s -> (s -> m b) -> Fold m a b
Fold forall {m :: * -> *} {a} {c}.
(MonadIO m, Storable a, Num c) =>
Tuple3' (Ring a) (Ptr a) c -> a -> m (Tuple3' (Ring a) (Ptr a) c)
step m (Tuple3' (Ring a) (Ptr a) Int)
initial forall {m :: * -> *} {a}.
(MonadIO m, Storable a) =>
Tuple3' (Ring a) (Ptr a) Int -> m (Array a)
done
  where
    step :: Tuple3' (Ring a) (Ptr a) c -> a -> m (Tuple3' (Ring a) (Ptr a) c)
step (Tuple3' Ring a
rb Ptr a
rh c
i) a
a = do
        Ptr a
rh1 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Ring a -> Ptr a -> a -> IO (Ptr a)
RB.unsafeInsert Ring a
rb Ptr a
rh a
a
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b c. a -> b -> c -> Tuple3' a b c
Tuple3' Ring a
rb Ptr a
rh1 (c
i forall a. Num a => a -> a -> a
+ c
1)
    initial :: m (Tuple3' (Ring a) (Ptr a) Int)
initial = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Ring a
a, Ptr a
b) -> forall a b c. a -> b -> c -> Tuple3' a b c
Tuple3' Ring a
a Ptr a
b (Int
0 :: Int)) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Int -> IO (Ring a, Ptr a)
RB.new Int
n
    done :: Tuple3' (Ring a) (Ptr a) Int -> m (Array a)
done (Tuple3' Ring a
rb Ptr a
rh Int
i) = do
        Array a
arr <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Int -> IO (Array a)
A.newArray Int
n
        forall {m :: * -> *} {a} {b}.
(MonadIO m, Storable a) =>
Int -> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
foldFunc Int
i Ptr a
rh forall {m :: * -> *} {a}.
(MonadIO m, Storable a) =>
Array a -> a -> m (Array a)
snoc' Array a
arr Ring a
rb
    snoc' :: Array a -> a -> m (Array a)
snoc' Array a
b a
a = forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Array a -> a -> IO (Array a)
A.unsafeSnoc Array a
b a
a
    foldFunc :: Int -> Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
foldFunc Int
i
        | Int
i forall a. Ord a => a -> a -> Bool
< Int
n = forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingM
        | Bool
otherwise = forall (m :: * -> *) a b.
(MonadIO m, Storable a) =>
Ptr a -> (b -> a -> m b) -> b -> Ring a -> m b
RB.unsafeFoldRingFullM

------------------------------------------------------------------------------
-- Time related
------------------------------------------------------------------------------

-- XXX using getTime in the loop can be pretty expensive especially for
-- computations where iterations are lightweight. We have the following
-- options:
--
-- 1) Run a timeout thread updating a flag asynchronously and check that
-- flag here, that way we can have a cheap termination check.
--
-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
--
-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
-- termination value of rdtsc in the beginning and then in each iteration just
-- get rdtsc and check if we should terminate.
--
data TakeByTime st s
    = TakeByTimeInit st
    | TakeByTimeCheck st s
    | TakeByTimeYield st s

{-# INLINE_NORMAL takeByTime #-}
takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
takeByTime :: forall (m :: * -> *) t a.
(MonadIO m, TimeUnit64 t) =>
t -> Stream m a -> Stream m a
takeByTime t
duration (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> TakeByTime s AbsTime -> m (Step (TakeByTime s AbsTime) a)
step (forall st s. st -> TakeByTime st s
TakeByTimeInit s
state1)
    where

    lim :: RelTime64
lim = forall a. TimeUnit64 a => a -> RelTime64
toRelTime64 t
duration

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> TakeByTime s AbsTime -> m (Step (TakeByTime s AbsTime) a)
step State Stream m a
_ (TakeByTimeInit s
_) | RelTime64
lim forall a. Eq a => a -> a -> Bool
== RelTime64
0 = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step State Stream m a
_ (TakeByTimeInit s
st) = do
        AbsTime
t0 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Clock -> IO AbsTime
getTime Clock
Monotonic
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall st s. st -> s -> TakeByTime st s
TakeByTimeYield s
st AbsTime
t0)
    step State Stream m a
_ (TakeByTimeCheck s
st AbsTime
t0) = do
        AbsTime
t <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Clock -> IO AbsTime
getTime Clock
Monotonic
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
            if AbsTime -> AbsTime -> RelTime64
diffAbsTime64 AbsTime
t AbsTime
t0 forall a. Ord a => a -> a -> Bool
> RelTime64
lim
            then forall s a. Step s a
Stop
            else forall s a. s -> Step s a
Skip (forall st s. st -> s -> TakeByTime st s
TakeByTimeYield s
st AbsTime
t0)
    step State Stream m a
gst (TakeByTimeYield s
st AbsTime
t0) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
             Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x (forall st s. st -> s -> TakeByTime st s
TakeByTimeCheck s
s AbsTime
t0)
             Skip s
s -> forall s a. s -> Step s a
Skip (forall st s. st -> s -> TakeByTime st s
TakeByTimeCheck s
s AbsTime
t0)
             Step s a
Stop -> forall s a. Step s a
Stop

data DropByTime st s x
    = DropByTimeInit st
    | DropByTimeGen st s
    | DropByTimeCheck st s x
    | DropByTimeYield st

{-# INLINE_NORMAL dropByTime #-}
dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
dropByTime :: forall (m :: * -> *) t a.
(MonadIO m, TimeUnit64 t) =>
t -> Stream m a -> Stream m a
dropByTime t
duration (Stream State Stream m a -> s -> m (Step s a)
step1 s
state1) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a
-> DropByTime s AbsTime a -> m (Step (DropByTime s AbsTime a) a)
step (forall st s x. st -> DropByTime st s x
DropByTimeInit s
state1)
    where

    lim :: RelTime64
lim = forall a. TimeUnit64 a => a -> RelTime64
toRelTime64 t
duration

    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> DropByTime s AbsTime a -> m (Step (DropByTime s AbsTime a) a)
step State Stream m a
_ (DropByTimeInit s
st) = do
        AbsTime
t0 <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Clock -> IO AbsTime
getTime Clock
Monotonic
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall st s x. st -> s -> DropByTime st s x
DropByTimeGen s
st AbsTime
t0)
    step State Stream m a
gst (DropByTimeGen s
st AbsTime
t0) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
             Yield a
x s
s -> forall s a. s -> Step s a
Skip (forall st s x. st -> s -> x -> DropByTime st s x
DropByTimeCheck s
s AbsTime
t0 a
x)
             Skip s
s -> forall s a. s -> Step s a
Skip (forall st s x. st -> s -> DropByTime st s x
DropByTimeGen s
s AbsTime
t0)
             Step s a
Stop -> forall s a. Step s a
Stop
    step State Stream m a
_ (DropByTimeCheck s
st AbsTime
t0 a
x) = do
        AbsTime
t <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ Clock -> IO AbsTime
getTime Clock
Monotonic
        if AbsTime -> AbsTime -> RelTime64
diffAbsTime64 AbsTime
t AbsTime
t0 forall a. Ord a => a -> a -> Bool
<= RelTime64
lim
        then forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall st s x. st -> s -> DropByTime st s x
DropByTimeGen s
st AbsTime
t0
        else forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x forall a b. (a -> b) -> a -> b
$ forall st s x. st -> DropByTime st s x
DropByTimeYield s
st
    step State Stream m a
gst (DropByTimeYield s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 State Stream m a
gst s
st
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s a
r of
             Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x (forall st s x. st -> DropByTime st s x
DropByTimeYield s
s)
             Skip s
s -> forall s a. s -> Step s a
Skip (forall st s x. st -> DropByTime st s x
DropByTimeYield s
s)
             Step s a
Stop -> forall s a. Step s a
Stop

-- XXX we should move this to stream generation section of this file. Also, the
-- take/drop combinators above should be moved to filtering section.
{-# INLINE_NORMAL currentTime #-}
currentTime :: MonadAsync m => Double -> Stream m AbsTime
currentTime :: forall (m :: * -> *). MonadAsync m => Double -> Stream m AbsTime
currentTime Double
g = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {p}.
(MonadIO m, MonadBaseControl IO m) =>
p
-> Maybe (Var IO Int64, ThreadId)
-> m (Step (Maybe (Var IO Int64, ThreadId)) AbsTime)
step forall a. Maybe a
Nothing

    where

    g' :: Double
g' = Double
g forall a. Num a => a -> a -> a
* Double
10 forall a b. (Num a, Integral b) => a -> b -> a
^ (Int
6 :: Int)

    -- XXX should have a minimum granularity to avoid high CPU usage?
    {-# INLINE delayTime #-}
    delayTime :: Int
delayTime =
        if Double
g' forall a. Ord a => a -> a -> Bool
>= forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Bounded a => a
maxBound :: Int)
        then forall a. Bounded a => a
maxBound
        else forall a b. (RealFrac a, Integral b) => a -> b
round Double
g'

    updateTimeVar :: Var IO Int64 -> IO ()
updateTimeVar Var IO Int64
timeVar = do
        Int -> IO ()
threadDelay forall a b. (a -> b) -> a -> b
$ Int
delayTime
        MicroSecond64 Int64
t <- forall a. TimeUnit a => AbsTime -> a
fromAbsTime forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Clock -> IO AbsTime
getTime Clock
Monotonic
        forall (m :: * -> *) a.
(MonadMut m, Prim a) =>
Var m a -> (a -> a) -> m ()
modifyVar' Var IO Int64
timeVar (forall a b. a -> b -> a
const Int64
t)

    {-# INLINE_LATE step #-}
    step :: p
-> Maybe (Var IO Int64, ThreadId)
-> m (Step (Maybe (Var IO Int64, ThreadId)) AbsTime)
step p
_ Maybe (Var IO Int64, ThreadId)
Nothing = do
        -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
        -- machine a 64-bit 'Var' cannot be read consistently without a lock
        -- while another thread is writing to it.
        Var IO Int64
timeVar <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadMut m, Prim a) => a -> m (Var m a)
newVar (Int64
0 :: Int64)
        ThreadId
tid <- forall (m :: * -> *).
(MonadIO m, MonadBaseControl IO m) =>
m () -> m ThreadId
forkManaged forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (Var IO Int64 -> IO ()
updateTimeVar Var IO Int64
timeVar)
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (Var IO Int64
timeVar, ThreadId
tid)

    step p
_ s :: Maybe (Var IO Int64, ThreadId)
s@(Just (Var IO Int64
timeVar, ThreadId
_)) = do
        Int64
a <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. (MonadMut m, Prim a) => Var m a -> m a
readVar Var IO Int64
timeVar
        -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
        -- efficiency.  or maybe we can use a representation using Double for
        -- floating precision time
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield (forall a. TimeUnit a => a -> AbsTime
toAbsTime (Int64 -> MicroSecond64
MicroSecond64 Int64
a)) Maybe (Var IO Int64, ThreadId)
s