Stream.hs

{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE UndecidableInstances #-}

module Data.Stream where

-- base
import Control.Applicative (Alternative (..), Applicative (..), liftA2)
import Control.Monad ((<$!>))
import Data.Bifunctor (bimap)
import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Monoid (Ap (..))
import Data.Tuple (swap)
import Prelude hiding (Applicative (..))

-- transformers
import Control.Monad.Trans.Class
import Control.Monad.Trans.Except (ExceptT (..), except, runExceptT, throwE, withExceptT)
import Control.Monad.Trans.Maybe (MaybeT (..))
import Control.Monad.Trans.Writer (WriterT (runWriterT), writer)

-- mmorph
import Control.Monad.Morph (MFunctor (hoist))

-- simple-affine-space
import Data.VectorSpace (VectorSpace (..))

-- selective
import Control.Selective

-- these
import Data.These (These (..))

-- semialign
import Data.Align

-- automaton
import Data.Stream.Internal
import Data.Stream.Recursive (Recursive (..))
import Data.Stream.Result

-- * Creating streams

{- | Effectful streams in coalgebraic encoding.

A stream consists of an internal state @s@, and a step function.
This step can make use of an effect in @m@ (which is often a monad),
alter the state, and return a result value.
Its semantics is continuously outputting values of type @b@,
while performing side effects in @m@.

A coalgebraic encoding was chosen instead of the direct recursion known from e.g. @list-transformer@, @dunai@, @machines@, @streaming@, ...,
because the coalgebraic encoding is much more amenable to compiler optimizations
than the coalgebraic encoding, which is:

@
  data StreamRecursiveT m b = StreamRecursiveT (m (b, StreamRecursiveT m b))
@

When two streams are composed, GHC can often optimize the combined step function,
resulting in a faster streams than what the coalgebraic encoding can ever achieve,
because the coalgebraic encoding has to step through every continuation.
Put differently, the compiler can perform static analysis on the state types of initially encoded state machines,
while the coalgebraic encoding knows its state only at runtime.

This performance gain comes at a peculiar cost:
Recursive definitions /of/ streams are not possible, e.g. an equation like:
@
  fixA stream = stream <*> fixA stream
@
This is impossible since the stream under definition itself appears in the definition body,
and thus the internal /state type/ would be recursively defined, which GHC doesn't allow:
Type level recursion is not supported in existential types.
An stream defined thusly will typically hang and/or leak memory, trying to build up an infinite type at runtime.

It is nevertheless possible to define streams recursively, but one needs to first identify the recursive definition of its /state type/.
Then for the greatest generality, 'fixStream' and 'fixStream'' can be used, and some special cases are covered by functions
such as 'fixA', 'Data.Automaton.parallely', 'many' and 'some'.
-}
data StreamT m a
  = forall s.
  StreamT
  { state :: s
  -- ^ The internal state of the stream
  , step :: s -> m (Result s a)
  -- ^ Stepping a stream by one tick means:
  --   1. performing a side effect in @m@
  --   2. updating the internal state @s@
  --   3. outputting a value of type @a@
  }

-- | Initialise with an internal state, update the state and produce output without side effects.
unfold :: (Applicative m) => s -> (s -> Result s a) -> StreamT m a
unfold state step =
  StreamT
    { state
    , step = pure . step
    }

-- | Like 'unfold', but output the current state.
unfold_ :: (Applicative m) => s -> (s -> s) -> StreamT m s
unfold_ state step = unfold state $ \s -> let s' = step s in Result s' s'

-- | Constantly perform the same effect, without remembering a state.
constM :: (Functor m) => m a -> StreamT m a
constM ma = StreamT () $ const $ Result () <$> ma
{-# INLINE constM #-}

-- | Like 'fmap' or 'rmap', but the postcomposed function may have an effect in @m@.
mmap :: (Monad m) => (a -> m b) -> StreamT m a -> StreamT m b
mmap f StreamT {state, step} =
  StreamT
    { state
    , step = \s -> do
        Result s' a <- step s
        Result s' <$> f a
    }
{-# INLINE mmap #-}

{- | Translate a coalgebraically encoded stream into a recursive one.

This is usually a performance penalty.
-}
toRecursive :: (Functor m) => StreamT m a -> Recursive m a
toRecursive automaton = Recursive $ mapResultState toRecursive <$> stepStream automaton
{-# INLINE toRecursive #-}

{- | Translate a recursive stream into a coalgebraically encoded one.

The internal state is the stream itself.
-}
fromRecursive :: Recursive m a -> StreamT m a
fromRecursive coalgebraic =
  StreamT
    { state = coalgebraic
    , step = getRecursive
    }
{-# INLINE fromRecursive #-}

-- | Call the monadic action once on the first tick and provide its result indefinitely.
initialised :: (Monad m) => m a -> StreamT m a
initialised action =
  let step mr@(Just r) = pure $! Result mr r
      step Nothing = (step . Just =<< action)
   in StreamT
        { state = Nothing
        , step
        }
{-# INLINE initialised #-}

instance (Functor m) => Functor (StreamT m) where
  fmap f StreamT {state, step} = StreamT state $! fmap (fmap f) <$> step
  {-# INLINE fmap #-}

-- | 'pure' forever returns the same value, '(<*>)' steps two streams synchronously.
instance (Applicative m) => Applicative (StreamT m) where
  pure = constM . pure
  {-# INLINE pure #-}

  StreamT stateF0 stepF <*> StreamT stateA0 stepA =
    StreamT (JointState stateF0 stateA0) (\(JointState stateF stateA) -> apResult <$> stepF stateF <*> stepA stateA)
  {-# INLINE (<*>) #-}

instance (Foldable m) => Foldable (StreamT m) where
  foldMap f StreamT {state, step} = go state
    where
      go s = step s & foldMap (\(Result s' a) -> f a <> go s')

instance (Traversable m, Functor m) => Traversable (StreamT m) where
  traverse f = fmap fromRecursive . traverse f . toRecursive

deriving via Ap (StreamT m) a instance (Applicative m, Num a) => Num (StreamT m a)

instance (Applicative m, Fractional a) => Fractional (StreamT m a) where
  fromRational = pure . fromRational
  recip = fmap recip

instance (Applicative m, Floating a) => Floating (StreamT m a) where
  pi = pure pi
  exp = fmap exp
  log = fmap log
  sin = fmap sin
  cos = fmap cos
  asin = fmap asin
  acos = fmap acos
  atan = fmap atan
  sinh = fmap sinh
  cosh = fmap cosh
  asinh = fmap asinh
  acosh = fmap acosh
  atanh = fmap atanh

instance (VectorSpace v s, Eq s, Floating s, Applicative m) => VectorSpace (StreamT m v) (StreamT m s) where
  zeroVector = pure zeroVector
  (*^) = liftA2 (*^)
  (^+^) = liftA2 (^+^)
  dot = liftA2 dot
  normalize = fmap normalize

{- | 'empty' just performs 'empty' in the underlying monad @m@.
  @s1 '<|>' s2@ starts in an undecided state,
  and explores the possibilities of continuing in @s1@ or @s2@
  on the first tick, using the underlying @m@.
-}
instance (Alternative m) => Alternative (StreamT m) where
  empty = constM empty
  {-# INLINE empty #-}

  StreamT stateL0 stepL <|> StreamT stateR0 stepR =
    StreamT
      { state = Undecided
      , step = \case
          Undecided -> (mapResultState DecideL <$> stepL stateL0) <|> (mapResultState DecideR <$> stepR stateR0)
          DecideL stateL -> mapResultState DecideL <$> stepL stateL
          DecideR stateR -> mapResultState DecideR <$> stepR stateR
      }
  {-# INLINE (<|>) #-}

  many StreamT {state, step} = fixStream'
    (const NotStarted)
    $ \fixstate fixstep -> \case
      NotStarted -> ((\(Result s' a) (Result ss' as) -> Result (Ongoing ss' s') $ a : as) <$> step state <*> fixstep fixstate) <|> pure (Result Finished [])
      Finished -> pure $! Result Finished []
      Ongoing ss s -> (\(Result s' a) (Result ss' as) -> Result (Ongoing ss' s') $ a : as) <$> step s <*> fixstep ss
  {-# INLINE many #-}

  some stream = (:) <$> stream <*> many stream
  {-# INLINE some #-}

instance MFunctor StreamT where
  hoist = hoist'
  {-# INLINE hoist #-}

{- | Hoist a stream along a monad morphism, by applying said morphism to the step function.

This is like @mmorph@'s 'hoist', but it doesn't require a 'Monad' constraint on @m2@.
-}
hoist' :: (forall x. m1 x -> m2 x) -> StreamT m1 a -> StreamT m2 a
hoist' f StreamT {state, step} = StreamT {state, step = f <$> step}
{-# INLINE hoist' #-}

-- * Running streams

-- | Perform one step of a stream, resulting in an updated stream and an output value.
stepStream :: (Functor m) => StreamT m a -> m (Result (StreamT m a) a)
stepStream StreamT {state, step} = mapResultState (`StreamT` step) <$> step state
{-# INLINE stepStream #-}

{- | Run a stream with trivial output.

If the output of a stream does not contain information,
all of its meaning is in its effects.
This function runs the stream indefinitely.
Since it will never return with a value, this function also has no output (its output is void).
The only way it can return is if @m@ includes some effect of termination,
e.g. 'Maybe' or 'Either' could terminate with a 'Nothing' or 'Left' value,
or 'IO' can raise an exception.
-}
reactimate :: (Monad m) => StreamT m () -> m void
reactimate StreamT {state, step} = go state
  where
    go s = do
      Result s' () <- step s
      go s'
{-# INLINE reactimate #-}

-- | Run a stream, collecting the outputs in a lazy, infinite list.
streamToList :: (Monad m) => StreamT m a -> m [a]
streamToList StreamT {state, step} = go state
  where
    go s = do
      Result s' a <- step s
      (a :) <$> go s'
{-# INLINE streamToList #-}

-- * Modifying streams

-- | Change the output type and effect of a stream without changing its state type.
withStreamT :: (Functor m, Functor n) => (forall s. m (Result s a) -> n (Result s b)) -> StreamT m a -> StreamT n b
withStreamT f StreamT {state, step} = StreamT state $ fmap f step
{-# INLINE withStreamT #-}

{- | Buffer the output of a stream, returning one value at a time.

This function lets a stream control the speed at which it produces data,
since it can decide to produce any amount of output at every step.
-}
concatS :: (Monad m) => StreamT m [a] -> StreamT m a
concatS StreamT {state, step} =
  StreamT
    { state = (state, [])
    , step = go
    }
  where
    go (s, []) = do
      Result s' as <- step s
      go (s', as)
    go (s, a : as) = pure $ Result (s, as) a
{-# INLINE concatS #-}

{- | At each step, duplicate the @m@ effect of the current step to the output.

This is useful if @m@ has some means of static analysis, or if you want to re-perform the effects.
-}
snapshot :: (Functor m) => StreamT m a -> StreamT m (m a)
snapshot StreamT {state, step} =
  StreamT
    { state
    , step = \s ->
        let result = step s
         in flip Result (output <$> result) . resultState <$> result
    }

-- ** Exception handling

{- | Streams with exceptions are 'Applicative' in the exception type.

Run the first stream until it throws a function as an exception,
then run the second one. If the second one ever throws an exception,
apply the function thrown by the first one to it.
-}
applyExcept :: (Monad m) => StreamT (ExceptT (e1 -> e2) m) a -> StreamT (ExceptT e1 m) a -> StreamT (ExceptT e2 m) a
applyExcept (StreamT state1 step1) (StreamT state2 step2) =
  StreamT
    { state = Left state1
    , step
    }
  where
    step (Left s1) = do
      resultOrException <- lift $ runExceptT $ step1 s1
      case resultOrException of
        Right result -> pure $! mapResultState Left result
        Left f -> step (Right (state2, f))
    step (Right (s2, f)) = mapResultState (Right . (,f)) <$!> withExceptT f (step2 s2)
{-# INLINE applyExcept #-}

{- | Execute the stream until it throws an exception, then restart it.

One might be tempted to define this function recursively with 'applyExcept',
but this would result in a runtime error, trying to define an infinite state.
-}
foreverExcept :: (Functor m, Monad m) => StreamT (ExceptT e m) a -> StreamT m a
foreverExcept StreamT {state, step} =
  StreamT
    { state
    , step = stepNew
    }
  where
    stepNew s = do
      resultOrException <- runExceptT $ step s
      case resultOrException of
        Left _ -> stepNew state
        Right result -> pure result

-- | Whenever an exception occurs, output it and retry on the next step.
exceptS :: (Applicative m) => StreamT (ExceptT e m) b -> StreamT m (Either e b)
exceptS StreamT {state, step} =
  StreamT
    { step = \state -> fmap (either (Result state . Left) (fmap Right)) $ runExceptT $ step state
    , state
    }
{-# INLINE exceptS #-}

{- | Run the first stream until it throws an exception.
  If the exception is 'Right', throw it immediately.
  If it is 'Left', run the second stream until it throws a function, which is then applied to the first exception.
-}
selectExcept :: (Monad m) => StreamT (ExceptT (Either e1 e2) m) a -> StreamT (ExceptT (e1 -> e2) m) a -> StreamT (ExceptT e2 m) a
selectExcept (StreamT stateE0 stepE) (StreamT stateF0 stepF) =
  StreamT
    { state = Left stateE0
    , step
    }
  where
    step (Left stateE) = do
      resultOrException <- lift $ runExceptT $ stepE stateE
      case resultOrException of
        Right result -> pure $ mapResultState Left result
        Left (Left e1) -> step (Right (e1, stateF0))
        Left (Right e2) -> throwE e2
    step (Right (e1, stateF)) = withExceptT ($ e1) $ mapResultState (Right . (e1,)) <$> stepF stateF

instance (Selective m) => Selective (StreamT m) where
  select (StreamT stateE0 stepE) (StreamT stateF0 stepF) =
    StreamT
      { state = JointState stateE0 stateF0
      , step = \(JointState stateE stateF) ->
          (fmap (mapResultState (`JointState` stateF)) . eitherResult <$> stepE stateE)
            <*? ((\(Result stateF' f) (Result stateE' a) -> Result (JointState stateE' stateF') (f a)) <$> stepF stateF)
      }
    where
      eitherResult :: Result s (Either a b) -> Either (Result s a) (Result s b)
      eitherResult (Result s eab) = bimap (Result s) (Result s) eab

{- | Run both streams in parallel and use @'Semialign' m@ to decide which stream produces output.
  If you understand @m@ as an effect that models the passage of time, then 'align' runs both streams concurrently.
-}
instance (Semialign m) => Semialign (StreamT m) where
  align (StreamT s10 step1) (StreamT s20 step2) =
    StreamT
      { state = JointState s10 s20
      , step = \(JointState s1 s2) -> align (step1 s1) (step2 s2) <&> updateTheseState s1 s2
      }
    where
      updateTheseState :: s1 -> s2 -> These (Result s1 a) (Result s2 b) -> Result (JointState s1 s2) (These a b)
      updateTheseState _s1 s2 (This (Result s1 a)) = Result (JointState s1 s2) $ This a
      updateTheseState s1 _s2 (That (Result s2 b)) = Result (JointState s1 s2) $ That b
      updateTheseState _ _ (These (Result s1 a) (Result s2 b)) = Result (JointState s1 s2) $ These a b
  {-# INLINE align #-}

instance (Align m) => Align (StreamT m) where
  nil = constM nil
  {-# INLINE nil #-}

-- ** Fix points, or recursive definitions

{- | Recursively define a stream from a recursive definition of the state, and of the step function.

If you want to define a stream recursively, this is not possible directly.
For example, consider this definition:
@
loops :: Monad m => StreamT m [Int]
loops = (:) <$> unfold_ 0 (+ 1) <*> loops
@
The defined value @loops@ contains itself in its definition.
This means that the internal state type of @loops@ must itself be recursively defined.
But GHC cannot do this automatically, because type level and value level are separate.
Instead, we need to spell out the type level recursion explicitly with a type constructor,
over which we will take the fixpoint.

In this example, we can figure out from the definitions that:
1. @'unfold_' 0 (+ 1)@ has @0 :: Int@ as state
2. '(:)' does not change the state
3. '<*>' takes the product of both states

So the internal state @s@ of @loops@ must satisfy the equation @s = (Int, s)@.
If the recursion is written as above, it tries to compute the infinite tuple @(Int, (Int, (Int, ...)))@, which hangs.
Instead, we need to define a type operator over which we take the fixpoint:

@
-- You need to write this:
data Loops x = Loops Int x

-- The library supplies:
data Fix f = Fix f (Fix f)
type LoopsState = Fix Loops
@

We can then use 'fixStream' to define the recursive definition of @loops@.
For this, we have to to tediously inline the definitions of 'unfold_', '(:)', and '<*>',
until we arrive at an explicit recursive definition of the state and the step function of @loops@, separately.
These are the two arguments of 'fixStream'.

@
loops :: Monad m => StreamT m [Int]
loops = fixStream (Loops 0) $ \fixStep (Loops n fixState) -> do
  Result s' a <- fixStep fixState
  return $ Result (Loops (n + 1) s') a
@
-}
fixStream ::
  (Functor m) =>
  -- | The recursive definition of the state of the stream.
  (forall s. s -> t s) ->
  -- | The recursive definition of the step function of the stream.
  ( forall s.
    (s -> m (Result s a)) ->
    (t s -> m (Result (t s) a))
  ) ->
  StreamT m a
fixStream transformState transformStep =
  StreamT
    { state = fixState transformState
    , step
    }
  where
    step Fix {getFix} = mapResultState Fix <$> transformStep step getFix

-- | A generalisation of 'fixStream' where the step definition is allowed to depend on the state.
fixStream' ::
  (Functor m) =>
  (forall s. s -> t s) ->
  -- | The recursive definition of the state of the stream.
  (forall s. s -> (s -> m (Result s a)) -> (t s -> m (Result (t s) a))) ->
  -- | The recursive definition of the step function of the stream.
  StreamT m a
fixStream' transformState transformStep =
  StreamT
    { state = fixState transformState
    , step
    }
  where
    step fix@(Fix {getFix}) = mapResultState Fix <$> transformStep fix step getFix

{- | The solution to the equation @'fixA' stream = stream <*> 'fixA' stream@.

Such a fix point operator needs to be used instead of the above direct definition because recursive definitions of streams
loop at runtime due to the coalgebraic encoding of the state.
-}
fixA :: (Applicative m) => StreamT m (a -> a) -> StreamT m a
fixA StreamT {state, step} = fixStream (JointState state) $
  \stepA (JointState s ss) -> apResult <$> step s <*> stepA ss

-- * Effect handling

-- | Lift the monad of a stream into a transformer.
liftS :: (Monad m, MonadTrans t) => StreamT m a -> StreamT (t m) a
liftS = hoist lift

{- | Continuously interpret a first order effect.

Several types are relevant here:

* @sig@: An effect signature functor, that encodes one effect.
  For example, @'Either' e@ for raising exceptions of type @e@, or @(w, )@ for a logging effect.
* @eff@: A monad that carries the effect.
  This can be a monad transformer stack including a transformer corresponding to @sig@, such as 'ExceptT' for 'Either'.
  It can also be the @Eff@ monad of an effect library such as @polysemy@, @bluefin@, @effectful@ and so on.
* @m@: The underlying monad in which the interpretation is performed, think "@eff@ without the effects from @sig@".

This function takes two functions, one to create effects in @eff@ from the signature, and the other to fully interpret them in @m@,
storing the complete effect information in @sig@ again.
It then executes the given automaton, extracting the effect by interpretation, and sending it back in.
The execution semantics is that of the monad @eff@, while the pure effect of the whole computation is returned in the output, encoded in @sig@.

For examples, see 'handleExceptT', 'handleWriterT' and similar functions below.
-}
handleEffect ::
  (Monad m, Monad eff, Functor sig) =>
  -- | Send a declarative effect in the signature to the effect carrier monad.
  (forall x. sig x -> eff x) ->
  -- | Interpret the effect in @m@, returning its result in the signature.
  (forall x. eff x -> m (sig x)) ->
  StreamT eff a ->
  StreamT m (sig a)
handleEffect send interpret StreamT {state, step} =
  StreamT
    { state = pure state
    , step = \s -> do
        results <- interpret $ step =<< s
        pure $! mapResultState send $ unzipResult results
    }

-- | Execute a stream until it throws an exception, then output the exception forever.
handleExceptT :: (Monad m) => StreamT (ExceptT e m) a -> StreamT m (Either e a)
handleExceptT = handleEffect except runExceptT

-- | Return the accumulated log at every step alongside the value.
handleWriterT :: (Monad m, Monoid w) => StreamT (WriterT w m) a -> StreamT m (w, a)
handleWriterT = handleEffect (writer . swap) (fmap swap . runWriterT)

-- | Execute a stream until it stops, then output 'Nothing' forever.
handleMaybeT :: (Monad m) => StreamT (MaybeT m) a -> StreamT m (Maybe a)
handleMaybeT = handleEffect (MaybeT . pure) runMaybeT