-- SPDX-FileCopyrightText: 2022 Gan Shen
-- SPDX-FileCopyrightText: 2025 Alex Ionescu
-- SPDX-License-Identifier: MPL-2.0 AND BSD-3-Clause

{-# LANGUAGE GADTs #-}

-- | This module defines the multi-thread backend for the `Network` monad.
module Choreography.Network.Local where

import Choreography.Location
import Choreography.Network
import Control.Concurrent
import Control.Monad
import Control.Monad.Freer
import Control.Monad.IO.Class
import Data.HashMap.Strict (HashMap, (!))
import Data.HashMap.Strict qualified as HashMap

-- | Each location is associated with a message buffer which stores messages sent
-- from other locations.
type MsgBuf = HashMap LocTm (Chan String)

newtype LocalConfig = LocalConfig
  { LocalConfig -> HashMap LocTm MsgBuf
locToBuf :: HashMap LocTm MsgBuf
  }

newEmptyMsgBuf :: [LocTm] -> IO MsgBuf
newEmptyMsgBuf :: [LocTm] -> IO MsgBuf
newEmptyMsgBuf = (MsgBuf -> LocTm -> IO MsgBuf) -> MsgBuf -> [LocTm] -> IO MsgBuf
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM MsgBuf -> LocTm -> IO MsgBuf
forall {k} {a}.
Hashable k =>
HashMap k (Chan a) -> k -> IO (HashMap k (Chan a))
f MsgBuf
forall k v. HashMap k v
HashMap.empty
  where
    f :: HashMap k (Chan a) -> k -> IO (HashMap k (Chan a))
f HashMap k (Chan a)
hash k
loc = do
      chan <- IO (Chan a)
forall a. IO (Chan a)
newChan
      return (HashMap.insert loc chan hash)

mkLocalConfig :: [LocTm] -> IO LocalConfig
mkLocalConfig :: [LocTm] -> IO LocalConfig
mkLocalConfig [LocTm]
locs = HashMap LocTm MsgBuf -> LocalConfig
LocalConfig (HashMap LocTm MsgBuf -> LocalConfig)
-> IO (HashMap LocTm MsgBuf) -> IO LocalConfig
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (HashMap LocTm MsgBuf -> LocTm -> IO (HashMap LocTm MsgBuf))
-> HashMap LocTm MsgBuf -> [LocTm] -> IO (HashMap LocTm MsgBuf)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM HashMap LocTm MsgBuf -> LocTm -> IO (HashMap LocTm MsgBuf)
f HashMap LocTm MsgBuf
forall k v. HashMap k v
HashMap.empty [LocTm]
locs
  where
    f :: HashMap LocTm MsgBuf -> LocTm -> IO (HashMap LocTm MsgBuf)
f HashMap LocTm MsgBuf
hash LocTm
loc = do
      buf <- [LocTm] -> IO MsgBuf
newEmptyMsgBuf [LocTm]
locs
      return (HashMap.insert loc buf hash)

locsLocal :: LocalConfig -> [LocTm]
locsLocal :: LocalConfig -> [LocTm]
locsLocal = HashMap LocTm MsgBuf -> [LocTm]
forall k v. HashMap k v -> [k]
HashMap.keys (HashMap LocTm MsgBuf -> [LocTm])
-> (LocalConfig -> HashMap LocTm MsgBuf) -> LocalConfig -> [LocTm]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LocalConfig -> HashMap LocTm MsgBuf
locToBuf

runNetworkLocal :: MonadIO m => LocalConfig -> LocTm -> Network m a -> m a
runNetworkLocal :: forall (m :: * -> *) a.
MonadIO m =>
LocalConfig -> LocTm -> Network m a -> m a
runNetworkLocal LocalConfig
cfg LocTm
self Network m a
prog = (forall a1. NetworkSig m a1 -> m a1) -> Network m a -> m a
forall (m :: * -> *) (f :: * -> *) a.
Monad m =>
(forall a1. f a1 -> m a1) -> Freer f a -> m a
interpFreer NetworkSig m a1 -> m a1
forall a1. NetworkSig m a1 -> m a1
forall (m :: * -> *) a. MonadIO m => NetworkSig m a -> m a
handler Network m a
prog
  where
    handler :: MonadIO m => NetworkSig m a -> m a
    handler :: forall (m :: * -> *) a. MonadIO m => NetworkSig m a -> m a
handler (Run m a
m)      = m a
m
    handler (Send a1
a LocTm
l)   = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ Chan LocTm -> LocTm -> IO ()
forall a. Chan a -> a -> IO ()
writeChan ((LocalConfig -> HashMap LocTm MsgBuf
locToBuf LocalConfig
cfg HashMap LocTm MsgBuf -> LocTm -> MsgBuf
forall k v.
(Eq k, Hashable k, HasCallStack) =>
HashMap k v -> k -> v
! LocTm
l) MsgBuf -> LocTm -> Chan LocTm
forall k v.
(Eq k, Hashable k, HasCallStack) =>
HashMap k v -> k -> v
! LocTm
self) (a1 -> LocTm
forall a. Show a => a -> LocTm
show a1
a)
    handler (Recv LocTm
l)     = IO a -> m a
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO a -> m a) -> IO a -> m a
forall a b. (a -> b) -> a -> b
$ LocTm -> a
forall a. Read a => LocTm -> a
read (LocTm -> a) -> IO LocTm -> IO a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Chan LocTm -> IO LocTm
forall a. Chan a -> IO a
readChan ((LocalConfig -> HashMap LocTm MsgBuf
locToBuf LocalConfig
cfg HashMap LocTm MsgBuf -> LocTm -> MsgBuf
forall k v.
(Eq k, Hashable k, HasCallStack) =>
HashMap k v -> k -> v
! LocTm
self) MsgBuf -> LocTm -> Chan LocTm
forall k v.
(Eq k, Hashable k, HasCallStack) =>
HashMap k v -> k -> v
! LocTm
l)
    handler (BCast a1
a [LocTm]
ls) = (NetworkSig m () -> m ()) -> [NetworkSig m ()] -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ NetworkSig m () -> m ()
forall (m :: * -> *) a. MonadIO m => NetworkSig m a -> m a
handler ([NetworkSig m ()] -> m ()) -> [NetworkSig m ()] -> m ()
forall a b. (a -> b) -> a -> b
$ (LocTm -> NetworkSig m ()) -> [LocTm] -> [NetworkSig m ()]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a1 -> LocTm -> NetworkSig m ()
forall a1 (m :: * -> *). Show a1 => a1 -> LocTm -> NetworkSig m ()
Send a1
a) [LocTm]
ls

instance Backend LocalConfig where
  locs :: LocalConfig -> [LocTm]
locs = LocalConfig -> [LocTm]
locsLocal
  runNetwork :: forall (m :: * -> *) a.
MonadIO m =>
LocalConfig -> LocTm -> Network m a -> m a
runNetwork = LocalConfig -> LocTm -> Network m a -> m a
forall (m :: * -> *) a.
MonadIO m =>
LocalConfig -> LocTm -> Network m a -> m a
runNetworkLocal