{-# LANGUAGE ExistentialQuantification #-}

module FRP.Rhine.Schedule.Internal where

-- base
import Control.Arrow
import Data.Function ((&))
import Data.Functor ((<&>))
import Data.Functor.Compose (Compose (..))
import Data.Kind (Type)
import Data.List.NonEmpty as N

-- foldable1-classes-compat
import Data.Foldable1 (Foldable1 (foldrMap1))

-- sop-core
import Data.SOP (HCollapse (hcollapse), HSequence (htraverse'), I (..), K (K), NP (..), NS (..), SListI, apInjs_NP, hliftA, hzipWith, unI)

-- monad-schedule
import Control.Monad.Schedule.Class

-- automaton
import Data.Stream hiding (concatS)
import Data.Stream.Result

-- | One step of a stream, with the state type argument going last, so it is usable with sop-core.
newtype Step m b state = Step {forall (m :: Type -> Type) b state.
Step m b state -> ResultStateT state m b
getStep :: ResultStateT state m b}

-- | The result of a stream, with the type arguments swapped, so it's usable with sop-core
newtype RunningResult b state = RunningResult {forall b state. RunningResult b state -> Result state b
getRunningResult :: Result state b}

-- | Transform an n-ary product of at least one type into a nonempty list of all its content.
apInjs_NPNonEmpty :: (SListI xs) => NP f (x ': xs) -> NonEmpty (NS f (x ': xs))
apInjs_NPNonEmpty :: forall (xs :: [Type]) (f :: Type -> Type) x.
SListI xs =>
NP f (x : xs) -> NonEmpty (NS f (x : xs))
apInjs_NPNonEmpty (f x
fx :* NP f xs
fxs) = f x -> NS f (x : xs)
forall {k} (a :: k -> Type) (x :: k) (xs :: [k]).
a x -> NS a (x : xs)
Z f x
f x
fx NS f (x : xs) -> [NS f (x : xs)] -> NonEmpty (NS f (x : xs))
forall a. a -> [a] -> NonEmpty a
:| (NS f xs -> NS f (x : xs)
forall {k} (a :: k -> Type) (xs :: [k]) (x :: k).
NS a xs -> NS a (x : xs)
S (NS f xs -> NS f (x : xs)) -> [NS f xs] -> [NS f (x : xs)]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> NP f xs -> [NS f xs]
forall {k} (xs :: [k]) (f :: k -> Type).
SListI xs =>
NP f xs -> [NS f xs]
apInjs_NP NP f xs
NP f xs
fxs)

-- | A nonempty list of 'StreamT's, unzipped into their states and their steps.
data Streams m b = forall state (states :: [Type]).
  (SListI states) =>
  Streams
  { ()
states :: NP I (state ': states)
  , ()
steps :: NP (Step m b) (state ': states)
  }

-- | Run 'Streams' concurrently by scheduling them in 'MonadSchedule'.
scheduleStreams :: (MonadSchedule m, Functor m, Applicative m) => Streams m b -> StreamT m (NonEmpty b)
scheduleStreams :: forall (m :: Type -> Type) b.
(MonadSchedule m, Functor m, Applicative m) =>
Streams m b -> StreamT m (NonEmpty b)
scheduleStreams Streams {NP I (state : states)
states :: ()
states :: NP I (state : states)
states, NP (Step m b) (state : states)
steps :: ()
steps :: NP (Step m b) (state : states)
steps} =
  StreamT
    { state :: (NonEmpty (NS I (state : states)),
 [m (NS (RunningResult b) (state : states))])
state = (NP I (state : states) -> NonEmpty (NS I (state : states))
forall (xs :: [Type]) (f :: Type -> Type) x.
SListI xs =>
NP f (x : xs) -> NonEmpty (NS f (x : xs))
apInjs_NPNonEmpty NP I (state : states)
states, []) -- All the initial states and no currently running continuations
    , step :: (NonEmpty (NS I (state : states)),
 [m (NS (RunningResult b) (state : states))])
-> m (Result
        (NonEmpty (NS I (state : states)),
         [m (NS (RunningResult b) (state : states))])
        (NonEmpty b))
step =
        -- Some streams have not started yet, or just finished their step. Others are still running.
        \(NonEmpty (NS I (state : states))
restingStates, [m (NS (RunningResult b) (state : states))]
runningStreams) ->
          -- Start all currently not running streams by zipping each with its step
          (NS I (state : states)
 -> m (NS (RunningResult b) (state : states)))
-> NonEmpty (NS I (state : states))
-> NonEmpty (m (NS (RunningResult b) (state : states)))
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap ((forall a. Compose m (RunningResult b) a -> m (RunningResult b a))
-> NS (Compose m (RunningResult b)) (state : states)
-> m (NS (RunningResult b) (state : states))
forall (xs :: [Type]) (g :: Type -> Type) (f :: Type -> Type)
       (f' :: Type -> Type).
(SListIN NS xs, Applicative g) =>
(forall a. f a -> g (f' a)) -> NS f xs -> g (NS f' xs)
forall k l (h :: (k -> Type) -> l -> Type) (xs :: l)
       (g :: Type -> Type) (f :: k -> Type) (f' :: k -> Type).
(HSequence h, SListIN h xs, Applicative g) =>
(forall (a :: k). f a -> g (f' a)) -> h f xs -> g (h f' xs)
htraverse' Compose m (RunningResult b) a -> m (RunningResult b a)
forall a. Compose m (RunningResult b) a -> m (RunningResult b a)
forall {k1} {k2} (f :: k1 -> Type) (g :: k2 -> k1) (a :: k2).
Compose f g a -> f (g a)
getCompose (NS (Compose m (RunningResult b)) (state : states)
 -> m (NS (RunningResult b) (state : states)))
-> (NS I (state : states)
    -> NS (Compose m (RunningResult b)) (state : states))
-> NS I (state : states)
-> m (NS (RunningResult b) (state : states))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall a. Step m b a -> I a -> Compose m (RunningResult b) a)
-> Prod NS (Step m b) (state : states)
-> NS I (state : states)
-> NS (Compose m (RunningResult b)) (state : states)
forall {k} {l} (h :: (k -> Type) -> l -> Type) (xs :: l)
       (f :: k -> Type) (f' :: k -> Type) (f'' :: k -> Type).
(SListIN (Prod h) xs, HAp h, HAp (Prod h)) =>
(forall (a :: k). f a -> f' a -> f'' a)
-> Prod h f xs -> h f' xs -> h f'' xs
hzipWith (\Step {ResultStateT a m b
getStep :: forall (m :: Type -> Type) b state.
Step m b state -> ResultStateT state m b
getStep :: ResultStateT a m b
getStep} -> m (RunningResult b a) -> Compose m (RunningResult b) a
forall {k} {k1} (f :: k -> Type) (g :: k1 -> k) (a :: k1).
f (g a) -> Compose f g a
Compose (m (RunningResult b a) -> Compose m (RunningResult b) a)
-> (I a -> m (RunningResult b a))
-> I a
-> Compose m (RunningResult b) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Result a b -> RunningResult b a)
-> m (Result a b) -> m (RunningResult b a)
forall a b. (a -> b) -> m a -> m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap Result a b -> RunningResult b a
forall b state. Result state b -> RunningResult b state
RunningResult (m (Result a b) -> m (RunningResult b a))
-> (I a -> m (Result a b)) -> I a -> m (RunningResult b a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ResultStateT a m b -> a -> m (Result a b)
forall s (m :: Type -> Type) a.
ResultStateT s m a -> s -> m (Result s a)
getResultStateT ResultStateT a m b
getStep (a -> m (Result a b)) -> (I a -> a) -> I a -> m (Result a b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. I a -> a
forall a. I a -> a
unI) Prod NS (Step m b) (state : states)
NP (Step m b) (state : states)
steps) NonEmpty (NS I (state : states))
restingStates
            -- Append all already running states to the freshly started ones
            NonEmpty (m (NS (RunningResult b) (state : states)))
-> (NonEmpty (m (NS (RunningResult b) (state : states)))
    -> NonEmpty (m (NS (RunningResult b) (state : states))))
-> NonEmpty (m (NS (RunningResult b) (state : states)))
forall a b. a -> (a -> b) -> b
& (NonEmpty (m (NS (RunningResult b) (state : states)))
 -> [m (NS (RunningResult b) (state : states))]
 -> NonEmpty (m (NS (RunningResult b) (state : states))))
-> [m (NS (RunningResult b) (state : states))]
-> NonEmpty (m (NS (RunningResult b) (state : states)))
-> NonEmpty (m (NS (RunningResult b) (state : states)))
forall a b c. (a -> b -> c) -> b -> a -> c
flip NonEmpty (m (NS (RunningResult b) (state : states)))
-> [m (NS (RunningResult b) (state : states))]
-> NonEmpty (m (NS (RunningResult b) (state : states)))
forall a. NonEmpty a -> [a] -> NonEmpty a
appendList [m (NS (RunningResult b) (state : states))]
runningStreams
            -- Schedule all running streams concurrently
            NonEmpty (m (NS (RunningResult b) (state : states)))
-> (NonEmpty (m (NS (RunningResult b) (state : states)))
    -> m (NonEmpty (NS (RunningResult b) (state : states)),
          [m (NS (RunningResult b) (state : states))]))
-> m (NonEmpty (NS (RunningResult b) (state : states)),
      [m (NS (RunningResult b) (state : states))])
forall a b. a -> (a -> b) -> b
& NonEmpty (m (NS (RunningResult b) (state : states)))
-> m (NonEmpty (NS (RunningResult b) (state : states)),
      [m (NS (RunningResult b) (state : states))])
forall a. NonEmpty (m a) -> m (NonEmpty a, [m a])
forall (m :: Type -> Type) a.
MonadSchedule m =>
NonEmpty (m a) -> m (NonEmpty a, [m a])
schedule
            -- Separate into finished streams and still running streams
            m (NonEmpty (NS (RunningResult b) (state : states)),
   [m (NS (RunningResult b) (state : states))])
-> (m (NonEmpty (NS (RunningResult b) (state : states)),
       [m (NS (RunningResult b) (state : states))])
    -> m (Result
            (NonEmpty (NS I (state : states)),
             [m (NS (RunningResult b) (state : states))])
            (NonEmpty b)))
-> m (Result
        (NonEmpty (NS I (state : states)),
         [m (NS (RunningResult b) (state : states))])
        (NonEmpty b))
forall a b. a -> (a -> b) -> b
& ((NonEmpty (NS (RunningResult b) (state : states)),
  [m (NS (RunningResult b) (state : states))])
 -> Result
      (NonEmpty (NS I (state : states)),
       [m (NS (RunningResult b) (state : states))])
      (NonEmpty b))
-> m (NonEmpty (NS (RunningResult b) (state : states)),
      [m (NS (RunningResult b) (state : states))])
-> m (Result
        (NonEmpty (NS I (state : states)),
         [m (NS (RunningResult b) (state : states))])
        (NonEmpty b))
forall a b. (a -> b) -> m a -> m b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap
              ( \(NonEmpty (NS (RunningResult b) (state : states))
finished, [m (NS (RunningResult b) (state : states))]
running) ->
                  let finishedStates :: NonEmpty (NS I (state : states))
finishedStates = NonEmpty (NS (RunningResult b) (state : states))
finished NonEmpty (NS (RunningResult b) (state : states))
-> (NS (RunningResult b) (state : states) -> NS I (state : states))
-> NonEmpty (NS I (state : states))
forall (f :: Type -> Type) a b. Functor f => f a -> (a -> b) -> f b
<&> ((forall a. RunningResult b a -> I a)
-> NS (RunningResult b) (state : states) -> NS I (state : states)
forall {k} {l} (h :: (k -> Type) -> l -> Type) (xs :: l)
       (f :: k -> Type) (f' :: k -> Type).
(SListIN (Prod h) xs, HAp h) =>
(forall (a :: k). f a -> f' a) -> h f xs -> h f' xs
hliftA (RunningResult b a -> Result a b
forall b state. RunningResult b state -> Result state b
getRunningResult (RunningResult b a -> Result a b)
-> (Result a b -> I a) -> RunningResult b a -> I a
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Result a b -> a
forall s a. Result s a -> s
resultState (Result a b -> a) -> (a -> I a) -> Result a b -> I a
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> a -> I a
forall a. a -> I a
I))
                      outputs :: NonEmpty b
outputs =
                        NonEmpty (NS (RunningResult b) (state : states))
finished
                          NonEmpty (NS (RunningResult b) (state : states))
-> (NS (RunningResult b) (state : states) -> b) -> NonEmpty b
forall (f :: Type -> Type) a b. Functor f => f a -> (a -> b) -> f b
<&> ((forall a. RunningResult b a -> K b a)
-> NS (RunningResult b) (state : states)
-> NS (K b) (state : states)
forall {k} {l} (h :: (k -> Type) -> l -> Type) (xs :: l)
       (f :: k -> Type) (f' :: k -> Type).
(SListIN (Prod h) xs, HAp h) =>
(forall (a :: k). f a -> f' a) -> h f xs -> h f' xs
hliftA (RunningResult b a -> Result a b
forall b state. RunningResult b state -> Result state b
getRunningResult (RunningResult b a -> Result a b)
-> (Result a b -> K b a) -> RunningResult b a -> K b a
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> Result a b -> b
forall s a. Result s a -> a
output (Result a b -> b) -> (b -> K b a) -> Result a b -> K b a
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> b -> K b a
forall k a (b :: k). a -> K a b
K) (NS (RunningResult b) (state : states)
 -> NS (K b) (state : states))
-> (NS (K b) (state : states) -> b)
-> NS (RunningResult b) (state : states)
-> b
forall {k} (cat :: k -> k -> Type) (a :: k) (b :: k) (c :: k).
Category cat =>
cat a b -> cat b c -> cat a c
>>> NS (K b) (state : states) -> b
NS (K b) (state : states) -> CollapseTo NS b
forall (xs :: [Type]) a.
SListIN NS xs =>
NS (K a) xs -> CollapseTo NS a
forall k l (h :: (k -> Type) -> l -> Type) (xs :: l) a.
(HCollapse h, SListIN h xs) =>
h (K a) xs -> CollapseTo h a
hcollapse)
                   in (NonEmpty (NS I (state : states)),
 [m (NS (RunningResult b) (state : states))])
-> NonEmpty b
-> Result
     (NonEmpty (NS I (state : states)),
      [m (NS (RunningResult b) (state : states))])
     (NonEmpty b)
forall s a. s -> a -> Result s a
Result (NonEmpty (NS I (state : states))
finishedStates, [m (NS (RunningResult b) (state : states))]
running) NonEmpty b
outputs
              )
    }

-- | Run a nonempty list of streams concurrently.
scheduleStreams' :: (MonadSchedule m, Applicative m) => NonEmpty (StreamT m b) -> StreamT m (NonEmpty b)
scheduleStreams' :: forall (m :: Type -> Type) b.
(MonadSchedule m, Applicative m) =>
NonEmpty (StreamT m b) -> StreamT m (NonEmpty b)
scheduleStreams' = Streams m b -> StreamT m (NonEmpty b)
forall (m :: Type -> Type) b.
(MonadSchedule m, Functor m, Applicative m) =>
Streams m b -> StreamT m (NonEmpty b)
scheduleStreams (Streams m b -> StreamT m (NonEmpty b))
-> (NonEmpty (StreamT m b) -> Streams m b)
-> NonEmpty (StreamT m b)
-> StreamT m (NonEmpty b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StreamT m b -> Streams m b)
-> (StreamT m b -> Streams m b -> Streams m b)
-> NonEmpty (StreamT m b)
-> Streams m b
forall a b. (a -> b) -> (a -> b -> b) -> NonEmpty a -> b
forall (t :: Type -> Type) a b.
Foldable1 t =>
(a -> b) -> (a -> b -> b) -> t a -> b
foldrMap1 StreamT m b -> Streams m b
forall (m :: Type -> Type) b. StreamT m b -> Streams m b
buildStreams StreamT m b -> Streams m b -> Streams m b
forall (m :: Type -> Type) b.
StreamT m b -> Streams m b -> Streams m b
consStreams
  where
    buildStreams :: StreamT m b -> Streams m b
    buildStreams :: forall (m :: Type -> Type) b. StreamT m b -> Streams m b
buildStreams StreamT {s
state :: ()
state :: s
state, s -> m (Result s b)
step :: ()
step :: s -> m (Result s b)
step} =
      Streams
        { states :: NP I '[s]
states = s -> I s
forall a. a -> I a
I s
state I s -> NP I '[] -> NP I '[s]
forall {k} (a :: k -> Type) (x :: k) (xs :: [k]).
a x -> NP a xs -> NP a (x : xs)
:* NP I '[]
forall {k} (a :: k -> Type). NP a '[]
Nil
        , steps :: NP (Step m b) '[s]
steps = ResultStateT s m b -> Step m b s
forall (m :: Type -> Type) b state.
ResultStateT state m b -> Step m b state
Step ((s -> m (Result s b)) -> ResultStateT s m b
forall s (m :: Type -> Type) a.
(s -> m (Result s a)) -> ResultStateT s m a
ResultStateT s -> m (Result s b)
step) Step m b s -> NP (Step m b) '[] -> NP (Step m b) '[s]
forall {k} (a :: k -> Type) (x :: k) (xs :: [k]).
a x -> NP a xs -> NP a (x : xs)
:* NP (Step m b) '[]
forall {k} (a :: k -> Type). NP a '[]
Nil
        }

    consStreams :: StreamT m b -> Streams m b -> Streams m b
    consStreams :: forall (m :: Type -> Type) b.
StreamT m b -> Streams m b -> Streams m b
consStreams StreamT {s
state :: ()
state :: s
state, s -> m (Result s b)
step :: ()
step :: s -> m (Result s b)
step} Streams {NP I (state : states)
states :: ()
states :: NP I (state : states)
states, NP (Step m b) (state : states)
steps :: ()
steps :: NP (Step m b) (state : states)
steps} =
      Streams
        { states :: NP I (s : state : states)
states = s -> I s
forall a. a -> I a
I s
state I s -> NP I (state : states) -> NP I (s : state : states)
forall {k} (a :: k -> Type) (x :: k) (xs :: [k]).
a x -> NP a xs -> NP a (x : xs)
:* NP I (state : states)
states
        , steps :: NP (Step m b) (s : state : states)
steps = ResultStateT s m b -> Step m b s
forall (m :: Type -> Type) b state.
ResultStateT state m b -> Step m b state
Step ((s -> m (Result s b)) -> ResultStateT s m b
forall s (m :: Type -> Type) a.
(s -> m (Result s a)) -> ResultStateT s m a
ResultStateT s -> m (Result s b)
step) Step m b s
-> NP (Step m b) (state : states)
-> NP (Step m b) (s : state : states)
forall {k} (a :: k -> Type) (x :: k) (xs :: [k]).
a x -> NP a xs -> NP a (x : xs)
:* NP (Step m b) (state : states)
steps
        }