{-# LANGUAGE Arrows #-}
{-# LANGUAGE CPP    #-}
-- |
-- Copyright  : (c) Ivan Perez and Manuel Baerenz, 2016
-- License    : BSD3
-- Maintainer : ivan.perez@keera.co.uk
--
-- Useful auxiliary functions and definitions.
module Data.MonadicStreamFunction.Util where

-- External imports
import Control.Arrow    (arr, returnA, (&&&), (<<<), (>>>))
import Control.Category (id, (.))
import Control.Monad    (when)
import Data.VectorSpace (VectorSpace, zeroVector, (^+^))
import Prelude          hiding (id, (.))

#if !MIN_VERSION_base(4,8,0)
import Data.Monoid (Monoid, mempty, mappend)
#endif

-- Internal imports
import Data.MonadicStreamFunction.Core                  (MSF, arrM, feedback)
import Data.MonadicStreamFunction.Instances.ArrowChoice ()

-- * Streams and sinks

-- | A stream is an 'MSF' that produces outputs, while ignoring the input. It
-- can obtain the values from a monadic context.
type MStream m a = MSF m () a

-- | A sink is an 'MSF' that consumes inputs, while producing no output. It
-- can consume the values with side effects.
type MSink m a = MSF m a ()

-- * Analogues of 'map' and 'fmap'

-- | Apply an 'MSF' to every input. Freezes temporarily if the input is
-- 'Nothing', and continues as soon as a 'Just' is received.
mapMaybeS :: Monad m => MSF m a b -> MSF m (Maybe a) (Maybe b)
mapMaybeS :: forall (m :: * -> *) a b.
Monad m =>
MSF m a b -> MSF m (Maybe a) (Maybe b)
mapMaybeS MSF m a b
msf = proc Maybe a
maybeA -> case Maybe a
maybeA of
  Just a
a  -> (b -> Maybe b) -> MSF m b (Maybe b)
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr b -> Maybe b
forall a. a -> Maybe a
Just MSF m b (Maybe b) -> MSF m a b -> MSF m a (Maybe b)
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
<<< MSF m a b
msf -< a
a
  Maybe a
Nothing -> MSF m (Maybe b) (Maybe b)
forall (a :: * -> * -> *) b. Arrow a => a b b
returnA          -< Maybe b
forall a. Maybe a
Nothing

-- * Adding side effects

-- | Applies a function to produce an additional side effect and passes the
-- input unchanged.
withSideEffect :: Monad m => (a -> m b) -> MSF m a a
withSideEffect :: forall (m :: * -> *) a b. Monad m => (a -> m b) -> MSF m a a
withSideEffect a -> m b
method = (MSF m a a
forall a. MSF m a a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
id MSF m a a -> MSF m a b -> MSF m a (a, b)
forall b c c'. MSF m b c -> MSF m b c' -> MSF m b (c, c')
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& (a -> m b) -> MSF m a b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> MSF m a b
arrM a -> m b
method) MSF m a (a, b) -> MSF m (a, b) a -> MSF m a a
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> ((a, b) -> a) -> MSF m (a, b) a
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (a, b) -> a
forall a b. (a, b) -> a
fst

-- | Produces an additional side effect and passes the input unchanged.
withSideEffect_ :: Monad m => m b -> MSF m a a
withSideEffect_ :: forall (m :: * -> *) b a. Monad m => m b -> MSF m a a
withSideEffect_ m b
method = (a -> m b) -> MSF m a a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> MSF m a a
withSideEffect ((a -> m b) -> MSF m a a) -> (a -> m b) -> MSF m a a
forall a b. (a -> b) -> a -> b
$ m b -> a -> m b
forall a b. a -> b -> a
const m b
method

-- * Delays

-- | Delay a signal by one sample.
iPre :: Monad m
     => a         -- ^ First output
     -> MSF m a a
iPre :: forall (m :: * -> *) a. Monad m => a -> MSF m a a
iPre a
firsta = a -> MSF m (a, a) (a, a) -> MSF m a a
forall (m :: * -> *) c a b.
Monad m =>
c -> MSF m (a, c) (b, c) -> MSF m a b
feedback a
firsta (MSF m (a, a) (a, a) -> MSF m a a)
-> MSF m (a, a) (a, a) -> MSF m a a
forall a b. (a -> b) -> a -> b
$ ((a, a) -> (a, a)) -> MSF m (a, a) (a, a)
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (a, a) -> (a, a)
forall {b} {a}. (b, a) -> (a, b)
swap
  where
    swap :: (b, a) -> (a, b)
swap (b
a, a
b) = (a
b, b
a)

-- | Preprends a fixed output to an 'MSF'. The first input is completely
-- ignored.
iPost :: Monad m => b -> MSF m a b -> MSF m a b
iPost :: forall (m :: * -> *) b a. Monad m => b -> MSF m a b -> MSF m a b
iPost b
b MSF m a b
sf = MSF m a b
sf MSF m a b -> MSF m b b -> MSF m a b
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (Maybe b -> MSF m (b, Maybe b) (b, Maybe b) -> MSF m b b
forall (m :: * -> *) c a b.
Monad m =>
c -> MSF m (a, c) (b, c) -> MSF m a b
feedback (b -> Maybe b
forall a. a -> Maybe a
Just b
b) (MSF m (b, Maybe b) (b, Maybe b) -> MSF m b b)
-> MSF m (b, Maybe b) (b, Maybe b) -> MSF m b b
forall a b. (a -> b) -> a -> b
$ ((b, Maybe b) -> (b, Maybe b)) -> MSF m (b, Maybe b) (b, Maybe b)
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (((b, Maybe b) -> (b, Maybe b)) -> MSF m (b, Maybe b) (b, Maybe b))
-> ((b, Maybe b) -> (b, Maybe b))
-> MSF m (b, Maybe b) (b, Maybe b)
forall a b. (a -> b) -> a -> b
$ \(b
c, Maybe b
ac) -> case Maybe b
ac of
  Maybe b
Nothing -> (b
c, Maybe b
forall a. Maybe a
Nothing)
  Just b
b' -> (b
b', Maybe b
forall a. Maybe a
Nothing))

-- | Preprends a fixed output to an 'MSF', shifting the output.
next :: Monad m => b -> MSF m a b -> MSF m a b
next :: forall (m :: * -> *) b a. Monad m => b -> MSF m a b -> MSF m a b
next b
b MSF m a b
sf = MSF m a b
sf MSF m a b -> MSF m b b -> MSF m a b
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> b -> MSF m b b
forall (m :: * -> *) a. Monad m => a -> MSF m a a
iPre b
b

-- | Buffers and returns the elements in FIFO order, returning 'Nothing'
-- whenever the buffer is empty.
fifo :: Monad m => MSF m [a] (Maybe a)
fifo :: forall (m :: * -> *) a. Monad m => MSF m [a] (Maybe a)
fifo = [a] -> MSF m ([a], [a]) (Maybe a, [a]) -> MSF m [a] (Maybe a)
forall (m :: * -> *) c a b.
Monad m =>
c -> MSF m (a, c) (b, c) -> MSF m a b
feedback [] ((([a], [a]) -> (Maybe a, [a])) -> MSF m ([a], [a]) (Maybe a, [a])
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr ([a] -> (Maybe a, [a])
forall x. [x] -> (Maybe x, [x])
safeSnoc ([a] -> (Maybe a, [a]))
-> (([a], [a]) -> [a]) -> ([a], [a]) -> (Maybe a, [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. ([a] -> [a] -> [a]) -> ([a], [a]) -> [a]
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry [a] -> [a] -> [a]
forall x. [x] -> [x] -> [x]
fifoAppend))
  where
    -- | Append a new list to an accumulator in FIFO order.
    fifoAppend :: [x] -> [x] -> [x]
    fifoAppend :: forall x. [x] -> [x] -> [x]
fifoAppend [x]
as [x]
accum = [x]
accum [x] -> [x] -> [x]
forall x. [x] -> [x] -> [x]
++ [x]
as

    -- | Split a list into the head and the tail.
    safeSnoc :: [x] -> (Maybe x, [x])
    safeSnoc :: forall x. [x] -> (Maybe x, [x])
safeSnoc []     = (Maybe x
forall a. Maybe a
Nothing, [])
    safeSnoc (x
x:[x]
xs) = (x -> Maybe x
forall a. a -> Maybe a
Just x
x, [x]
xs)

-- * Folding

-- ** Folding for 'VectorSpace' instances

-- | Count the number of simulation steps. Produces 1, 2, 3,...
count :: (Num n, Monad m) => MSF m a n
count :: forall n (m :: * -> *) a. (Num n, Monad m) => MSF m a n
count = (a -> n) -> MSF m a n
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (n -> a -> n
forall a b. a -> b -> a
const n
1) MSF m a n -> MSF m n n -> MSF m a n
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> (n -> n -> n) -> n -> MSF m n n
forall (m :: * -> *) a s.
Monad m =>
(a -> s -> s) -> s -> MSF m a s
accumulateWith n -> n -> n
forall a. Num a => a -> a -> a
(+) n
0

-- | Sums the inputs, starting from zero.
sumS :: (VectorSpace v s, Monad m) => MSF m v v
sumS :: forall v s (m :: * -> *). (VectorSpace v s, Monad m) => MSF m v v
sumS = v -> MSF m v v
forall v s (m :: * -> *).
(VectorSpace v s, Monad m) =>
v -> MSF m v v
sumFrom v
forall v a. VectorSpace v a => v
zeroVector

-- | Sums the inputs, starting from an initial vector.
sumFrom :: (VectorSpace v s, Monad m) => v -> MSF m v v
sumFrom :: forall v s (m :: * -> *).
(VectorSpace v s, Monad m) =>
v -> MSF m v v
sumFrom = (v -> v -> v) -> v -> MSF m v v
forall (m :: * -> *) a s.
Monad m =>
(a -> s -> s) -> s -> MSF m a s
accumulateWith v -> v -> v
forall v a. VectorSpace v a => v -> v -> v
(^+^)

-- ** Folding for monoids

-- | Accumulate the inputs, starting from 'mempty'.
mappendS :: (Monoid n, Monad m) => MSF m n n
mappendS :: forall n (m :: * -> *). (Monoid n, Monad m) => MSF m n n
mappendS = n -> MSF m n n
forall n (m :: * -> *). (Monoid n, Monad m) => n -> MSF m n n
mappendFrom n
forall a. Monoid a => a
mempty
{-# INLINE mappendS #-}

-- | Accumulate the inputs, starting from an initial monoid value.
mappendFrom :: (Monoid n, Monad m) => n -> MSF m n n
mappendFrom :: forall n (m :: * -> *). (Monoid n, Monad m) => n -> MSF m n n
mappendFrom = (n -> n -> n) -> n -> MSF m n n
forall (m :: * -> *) a s.
Monad m =>
(a -> s -> s) -> s -> MSF m a s
accumulateWith n -> n -> n
forall a. Monoid a => a -> a -> a
mappend

-- ** Generic folding \/ accumulation

-- | Applies a function to the input and an accumulator, outputting the updated
-- accumulator. Equal to @\f s0 -> feedback s0 $ arr (uncurry f >>> dup)@.
accumulateWith :: Monad m => (a -> s -> s) -> s -> MSF m a s
accumulateWith :: forall (m :: * -> *) a s.
Monad m =>
(a -> s -> s) -> s -> MSF m a s
accumulateWith a -> s -> s
f s
s0 = s -> MSF m (a, s) (s, s) -> MSF m a s
forall (m :: * -> *) c a b.
Monad m =>
c -> MSF m (a, c) (b, c) -> MSF m a b
feedback s
s0 (MSF m (a, s) (s, s) -> MSF m a s)
-> MSF m (a, s) (s, s) -> MSF m a s
forall a b. (a -> b) -> a -> b
$ ((a, s) -> (s, s)) -> MSF m (a, s) (s, s)
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (a, s) -> (s, s)
g
  where
    g :: (a, s) -> (s, s)
g (a
a, s
s) = let s' :: s
s' = a -> s -> s
f a
a s
s in (s
s', s
s')

-- | Applies a transfer function to the input and an accumulator, returning the
-- updated accumulator and output.
mealy :: Monad m => (a -> s -> (b, s)) -> s -> MSF m a b
mealy :: forall (m :: * -> *) a s b.
Monad m =>
(a -> s -> (b, s)) -> s -> MSF m a b
mealy a -> s -> (b, s)
f s
s0 = s -> MSF m (a, s) (b, s) -> MSF m a b
forall (m :: * -> *) c a b.
Monad m =>
c -> MSF m (a, c) (b, c) -> MSF m a b
feedback s
s0 (MSF m (a, s) (b, s) -> MSF m a b)
-> MSF m (a, s) (b, s) -> MSF m a b
forall a b. (a -> b) -> a -> b
$ ((a, s) -> (b, s)) -> MSF m (a, s) (b, s)
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (((a, s) -> (b, s)) -> MSF m (a, s) (b, s))
-> ((a, s) -> (b, s)) -> MSF m (a, s) (b, s)
forall a b. (a -> b) -> a -> b
$ (a -> s -> (b, s)) -> (a, s) -> (b, s)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> s -> (b, s)
f

-- * Unfolding

-- | Generate outputs using a step-wise generation function and an initial
-- value.
unfold :: Monad m => (a -> (b, a)) -> a -> MSF m () b
unfold :: forall (m :: * -> *) a b.
Monad m =>
(a -> (b, a)) -> a -> MSF m () b
unfold a -> (b, a)
f a
a = a -> MSF m ((), a) (b, a) -> MSF m () b
forall (m :: * -> *) c a b.
Monad m =>
c -> MSF m (a, c) (b, c) -> MSF m a b
feedback a
a ((((), a) -> (b, a)) -> MSF m ((), a) (b, a)
forall b c. (b -> c) -> MSF m b c
forall (a :: * -> * -> *) b c. Arrow a => (b -> c) -> a b c
arr (((), a) -> a
forall a b. (a, b) -> b
snd (((), a) -> a) -> (a -> (b, a)) -> ((), a) -> (b, a)
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> a -> (b, a)
f))

-- | Generate outputs using a step-wise generation function and an initial
-- value. Version of 'unfold' in which the output and the new accumulator are
-- the same. Should be equal to @\f a -> unfold (f >>> dup) a@.
repeatedly :: Monad m => (a -> a) -> a -> MSF m () a
repeatedly :: forall (m :: * -> *) a. Monad m => (a -> a) -> a -> MSF m () a
repeatedly a -> a
f = (a -> (a, a)) -> a -> MSF m () a
forall (m :: * -> *) a b.
Monad m =>
(a -> (b, a)) -> a -> MSF m () b
unfold ((a -> (a, a)) -> a -> MSF m () a)
-> (a -> (a, a)) -> a -> MSF m () a
forall a b. (a -> b) -> a -> b
$ a -> a
f (a -> a) -> (a -> (a, a)) -> a -> (a, a)
forall {k} (cat :: k -> k -> *) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> a -> (a, a)
forall {b}. b -> (b, b)
dup
  where
    dup :: b -> (b, b)
dup b
a = (b
a, b
a)

-- * Debugging

-- | Outputs every input sample, with a given message prefix.
trace :: Show a => String -> MSF IO a a
trace :: forall a. Show a => String -> MSF IO a a
trace = (String -> IO ()) -> String -> MSF IO a a
forall (m :: * -> *) a.
(Monad m, Show a) =>
(String -> m ()) -> String -> MSF m a a
traceWith String -> IO ()
putStrLn

-- | Outputs every input sample, with a given message prefix, using an
-- auxiliary printing function.
traceWith :: (Monad m, Show a) => (String -> m ()) -> String -> MSF m a a
traceWith :: forall (m :: * -> *) a.
(Monad m, Show a) =>
(String -> m ()) -> String -> MSF m a a
traceWith String -> m ()
method String
msg =
  (a -> m ()) -> MSF m a a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> MSF m a a
withSideEffect (String -> m ()
method (String -> m ()) -> (a -> String) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. (String
msg String -> String -> String
forall x. [x] -> [x] -> [x]
++) (String -> String) -> (a -> String) -> a -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. a -> String
forall a. Show a => a -> String
show)

-- | Outputs every input sample, with a given message prefix, using an
-- auxiliary printing function, when a condition is met.
traceWhen :: (Monad m, Show a)
          => (a -> Bool)
          -> (String -> m ())
          -> String
          -> MSF m a a
traceWhen :: forall (m :: * -> *) a.
(Monad m, Show a) =>
(a -> Bool) -> (String -> m ()) -> String -> MSF m a a
traceWhen a -> Bool
cond String -> m ()
method String
msg = (a -> m ()) -> MSF m a a
forall (m :: * -> *) a b. Monad m => (a -> m b) -> MSF m a a
withSideEffect ((a -> m ()) -> MSF m a a) -> (a -> m ()) -> MSF m a a
forall a b. (a -> b) -> a -> b
$ \a
a ->
  Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (a -> Bool
cond a
a) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ String -> m ()
method (String -> m ()) -> String -> m ()
forall a b. (a -> b) -> a -> b
$ String
msg String -> String -> String
forall x. [x] -> [x] -> [x]
++ a -> String
forall a. Show a => a -> String
show a
a

-- | Outputs every input sample, with a given message prefix, when a condition
-- is met, and waits for some input \/ enter to continue.
pauseOn :: Show a => (a -> Bool) -> String -> MSF IO a a
pauseOn :: forall a. Show a => (a -> Bool) -> String -> MSF IO a a
pauseOn a -> Bool
cond = (a -> Bool) -> (String -> IO ()) -> String -> MSF IO a a
forall (m :: * -> *) a.
(Monad m, Show a) =>
(a -> Bool) -> (String -> m ()) -> String -> MSF m a a
traceWhen a -> Bool
cond ((String -> IO ()) -> String -> MSF IO a a)
-> (String -> IO ()) -> String -> MSF IO a a
forall a b. (a -> b) -> a -> b
$ \String
s -> String -> IO ()
forall a. Show a => a -> IO ()
print String
s IO () -> IO String -> IO String
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> IO String
getLine IO String -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()