module Data.Machine.Fanout (fanout, fanoutSteps) where
import Control.Applicative
import Control.Arrow
import Control.Monad (foldM)
import Data.Machine
import Data.Maybe (catMaybes)
import Data.Monoid
import Data.Semigroup (Semigroup(sconcat))
import Data.List.NonEmpty (NonEmpty((:|)))
import Prelude
feed :: Monad m => a -> ProcessT m a b -> m (Step (Is a) b (ProcessT m a b))
feed x m = runMachineT m >>= \v ->
case v of
Await f Refl _ -> runMachineT (f x)
s -> return s
mapAccumLM :: (Functor m, Monad m)
=> (acc -> x -> m (acc, y)) -> acc -> [x] -> m (acc, [y])
mapAccumLM f z = fmap (second ($ [])) . foldM aux (z,id)
where aux (acc,ys) x = second ((. ys) . (:)) <$> f acc x
flushYields :: Monad m
=> Step k o (MachineT m k o) -> m ([o], Maybe (MachineT m k o))
flushYields = go id
where go rs (Yield o s) = runMachineT s >>= go ((o:) . rs)
go rs Stop = return (rs [], Nothing)
go rs s = return (rs [], Just $ encased s)
fanout :: (Functor m, Monad m, Semigroup r)
=> [ProcessT m a r] -> ProcessT m a r
fanout xs = encased $ Await (MachineT . aux) Refl (fanout xs)
where aux y = do (rs,xs') <- mapM (feed y) xs >>= mapAccumLM yields []
let nxt = fanout $ catMaybes xs'
case rs of
[] -> runMachineT nxt
(r:rs') -> return $ Yield (sconcat $ r :| rs') nxt
yields rs Stop = return (rs,Nothing)
yields rs y@(Yield _ _) = first (++ rs) <$> flushYields y
yields rs a@(Await _ _ _) = return (rs, Just $ encased a)
fanoutSteps :: (Functor m, Monad m, Monoid r)
=> [ProcessT m a r] -> ProcessT m a r
fanoutSteps xs = encased $ Await (MachineT . aux) Refl (fanoutSteps xs)
where aux y = do (rs,xs') <- mapM (feed y) xs >>= mapAccumLM yields []
let nxt = fanoutSteps $ catMaybes xs'
if null rs
then return $ Yield mempty nxt
else return $ Yield (mconcat rs) nxt
yields rs Stop = return (rs,Nothing)
yields rs y@(Yield _ _) = first (++rs) <$> flushYields y
yields rs a@(Await _ _ _) = return (rs, Just $ encased a)