Stream.hs

{-# LANGUAGE UndecidableInstances #-}

-- SPDX-License-Identifier: MPL-2.0

{- |
Copyright   :  (c) 2024 Sayo contributors
License     :  MPL-2.0 (see the LICENSE file)
Maintainer  :  ymdfield@outlook.jp

Coroutine-based, composable, and resumable concurrent streams.
-}
module Control.Monad.Hefty.Concurrent.Stream (
    module Control.Monad.Hefty.Concurrent.Stream,
    module Control.Monad.Hefty.Input,
    module Control.Monad.Hefty.Output,
)
where

import Control.Arrow (Arrow, ArrowChoice, arr, first, left, (>>>))
import Control.Category (Category)
import Control.Category qualified as C
import Control.Monad (forM_, forever)
import Control.Monad.Hefty (
    Eff,
    Emb,
    FOEs,
    RemoveHOEs,
    WeakenHOEs,
    interpret,
    interpretsBy,
    nil,
    onlyFOEs,
    raise,
    reinterprets,
    untag,
    (!:),
    (&),
    (:>),
    type (~>),
 )
import Control.Monad.Hefty.Concurrent.Parallel (Parallel, liftP2)
import Control.Monad.Hefty.Input
import Control.Monad.Hefty.Output
import Control.Monad.Hefty.State (State, evalState, evalStateIORef, get'', put'')
import Data.Effect.Unlift (UnliftIO, withRunInIO)
import Data.Function (fix)
import Data.Sequence (Seq ((:|>)))
import Data.Sequence qualified as Seq
import UnliftIO (
    atomically,
    liftIO,
    mask,
    newEmptyTMVarIO,
    putTMVar,
    readTMVar,
    takeTMVar,
    uninterruptibleMask_,
 )
import UnliftIO.Concurrent (forkIO, killThread)

data Machinery es ans i o where
    Unit
        :: forall i o ans es
         . Eff (Input i ': Output o ': es) ans
        -> Machinery es ans i o
    Connect
        :: forall a b c ans es
         . Machinery es ans a b
        -> Machinery es ans b c
        -> Machinery es ans a c

instance Category (Machinery es ans) where
    id :: forall a. Machinery es ans a a
    id =
        Unit . forever $
            input @a >>= output

    (.) = flip Connect

    {-# INLINE id #-}
    {-# INLINE (.) #-}

instance (FOEs es) => Arrow (Machinery es ans) where
    arr (f :: b -> c) =
        Unit . forever $
            input @b >>= output . f

    first
        :: forall b c d
         . Machinery es ans b c
        -> Machinery es ans (b, d) (c, d)
    first = \case
        Unit m -> Unit $ evalState (Left Seq.Empty) $ buffering m
        Connect a b -> Connect (first a) (first b)

    {-# INLINE arr #-}
    {-# INLINE first #-}

buffering
    :: forall b c d ans es
     . Eff (Input b ': Output c ': es) ans
    -> Eff (State (Either (Seq c) d) ': Input (b, d) ': Output (c, d) ': es) ans
buffering =
    reinterprets
        ( ( \Input -> do
                (b, d) <- input

                get'' @"buffer" >>= \case
                    Right _ -> pure ()
                    Left outputQueue -> forM_ outputQueue \c -> output (c, d)

                put'' @"buffer" $ Right d

                pure b
          )
            !: ( \(Output c) ->
                    get'' @"buffer" >>= \case
                        Right d -> output (c, d)
                        Left outputQueue -> put'' @"buffer" $ Left $ outputQueue :|> c
               )
            !: nil
        )
        >>> untag @"buffer"

instance (FOEs es) => ArrowChoice (Machinery es ans) where
    left = leftMachinery
    {-# INLINE left #-}

leftMachinery
    :: forall b c d ans es
     . Machinery es ans b c
    -> Machinery es ans (Either b d) (Either c d)
leftMachinery = \case
    Unit m ->
        m
            & reinterprets
                ( ( \Input -> fix \next ->
                        input @(Either b d) >>= \case
                            Left x -> pure x
                            Right o -> do
                                output @(Either c d) $ Right o
                                next
                  )
                    !: (\(Output o) -> output @(Either c d) $ Left o)
                    !: nil
                )
            & Unit
    Connect a b -> Connect (leftMachinery a) (leftMachinery b)

newtype Machine f ans i o = Machine
    {runMachine :: f (MachineStatus f ans i o)}

data MachineStatus f ans i o
    = Terminated ans
    | Waiting (i -> Machine f ans i o)
    | Produced o (Machine f ans i o)

machine :: (WeakenHOEs es) => Eff (Input i ': Output o ': RemoveHOEs es) ans -> Machine (Eff es) ans i o
machine =
    interpretsBy
        (pure . Terminated)
        ( (\Input k -> pure $ Waiting $ Machine . onlyFOEs . k)
            !: (\(Output o) k -> pure $ Produced o $ Machine $ onlyFOEs $ k ())
            !: nil
        )
        >>> onlyFOEs
        >>> Machine

runMachinery
    :: forall i o ans es
     . (Parallel :> es, Semigroup ans, WeakenHOEs es)
    => Machinery (RemoveHOEs es) ans i o
    -> Eff es (MachineStatus (Eff es) ans i o)
runMachinery = runMachineryL . mviewl

runMachineryL
    :: forall i o ans es
     . (Parallel :> es, Semigroup ans, WeakenHOEs es)
    => MachineryViewL (RemoveHOEs es) ans i o
    -> Eff es (MachineStatus (Eff es) ans i o)
runMachineryL = \case
    MOne m -> runMachine $ machine m
    MCons m ms -> do
        liftP2 (,) (runMachine $ machine m) (runMachinery ms) >>= loop
      where
        loop = \case
            (Terminated ans, Terminated ans') -> pure $ Terminated $ ans <> ans'
            (Produced o k1, Waiting k2) ->
                liftP2 (,) (runMachine k1) (runMachine $ k2 o) >>= loop
            (Waiting k, s) ->
                pure $ Waiting \i -> Machine do
                    s' <- runMachine $ k i
                    loop (s', s)
            (s, Produced o k) ->
                pure $ Produced o $ Machine do
                    s' <- runMachine k
                    loop (s, s')
            (Terminated ans, Waiting _) -> pure $ Terminated ans
            (Produced _ _, Terminated ans) -> pure $ Terminated ans

newtype MachineryIO es ans i o = MachineryIO {unMachineryIO :: Machinery es ans i o}
    deriving newtype (Category)

instance (Emb IO :> es) => Arrow (MachineryIO es ans) where
    arr (f :: b -> c) =
        MachineryIO . Unit . forever $
            input @b >>= output . f

    first :: forall b c d. MachineryIO es ans b c -> MachineryIO es ans (b, d) (c, d)
    first =
        unMachineryIO
            >>> MachineryIO . \case
                Unit m ->
                    Unit $ evalStateIORef (Left Seq.Empty) $ buffering m
                Connect a b ->
                    Connect
                        (unMachineryIO $ first $ MachineryIO a)
                        (unMachineryIO $ first $ MachineryIO b)

    {-# INLINE arr #-}
    {-# INLINE first #-}

instance (Emb IO :> es) => ArrowChoice (MachineryIO es ans) where
    left = MachineryIO . leftMachinery . unMachineryIO
    {-# INLINE left #-}

runMachineryIO
    :: forall i o ans es
     . (UnliftIO :> es, Emb IO :> es)
    => Eff es i
    -> (o -> Eff es ())
    -> Machinery es ans i o
    -> Eff es ans
runMachineryIO i o = runMachineryIOL i o . mviewl

runMachineryIOL
    :: forall i o ans es
     . (UnliftIO :> es, Emb IO :> es)
    => Eff es i
    -> (o -> Eff es ())
    -> MachineryViewL es ans i o
    -> Eff es ans
runMachineryIOL i o = \case
    MOne m -> runUnit o m
    MCons a b ->
        withRunInIO \run -> do
            chan <- newEmptyTMVarIO
            ans <- newEmptyTMVarIO
            mask \restore -> do
                let runThread m = forkIO do
                        x <- restore $ run m
                        atomically $ putTMVar ans x

                t1 <- runThread $ runUnit (liftIO . atomically . putTMVar chan) a
                t2 <- runThread $ runMachineryIO (liftIO . atomically $ takeTMVar chan) o b

                atomically (readTMVar ans)
                    <* uninterruptibleMask_ (killThread t1 *> killThread t2)
  where
    runUnit :: (o' -> Eff es ()) -> Eff (Input i ': Output o' ': es) ~> Eff es
    runUnit o' m =
        m
            & interpret (\Input -> raise i)
            & interpret (\(Output x) -> o' x)

runMachineryIO_
    :: forall ans es
     . (UnliftIO :> es, Emb IO :> es)
    => Machinery es ans () ()
    -> Eff es ans
runMachineryIO_ = runMachineryIO (pure ()) (const $ pure ())
{-# INLINE runMachineryIO_ #-}

-- Inspired by https://hackage.haskell.org/package/freer-simple-1.2.1.2/docs/Data-FTCQueue.html

{- |
Left view deconstruction data structure for Machinery Pipeline.

This allows the number of generated threads to be reduced to the number of machine units.
-}
data MachineryViewL es ans i o where
    MOne
        :: forall i o ans es
         . Eff (Input i ': Output o ': es) ans
        -> MachineryViewL es ans i o
    MCons
        :: forall a b c ans es
         . Eff (Input a ': Output b ': es) ans
        -> Machinery es ans b c
        -> MachineryViewL es ans a c

-- | Left view deconstruction for Machinery Pipeline. [average O(1)]
mviewl :: Machinery es ans i o -> MachineryViewL es ans i o
mviewl = \case
    Unit m -> MOne m
    Connect a b -> connect a b
  where
    connect
        :: Machinery es ans a b
        -> Machinery es ans b c
        -> MachineryViewL es ans a c
    connect (Unit m) r = m `MCons` r
    connect (Connect a b) r = connect a (Connect b r)