feat: federated learning demo on the actor framework

- supervisor: ignore stale death messages (OneForAll restart storm),
  drop normally-exited children instead of respawning, add a
  restart-rate limit
- register actors before their thread starts so UUID dispatch and
  getSelf cannot race the spawn
- killActor now kills the thread so links see Killed
- fulfill connection promises with the pool winner on simultaneous
  connects
- clamp the PaySim test split on small datasets and split shared
  transaction types across trainers in non-IID partitioning

.gitignore

@@ -1 +1,3 @@
 dist-newstyle
+demo-out/
+fl_model_*.bin

actors.cabal

@@ -62,7 +62,28 @@ library
     import:           warnings
 
     -- Modules exported by the library.
-    exposed-modules:  Control.Actor
+    exposed-modules:
+      Control.Actor
+      Control.Actor.Types
+      Control.Actor.Transport
+      Control.Actor.Runtime
+      Control.Actor.Core
+      Control.Actor.Supervision
+      Control.Actor.Network
+      Control.Actor.Registry
+      FL.Types
+      FL.CRDT
+      FL.NN
+      FL.Retry
+      FL.Data.PaySim
+      FL.Actors.Logger
+      FL.Actors.Evaluator
+      FL.Actors.Trainer
+      FL.Actors.Aggregator
+      FL.Actors.Coordinator
+      FL.Actors.PeerCoordinator
+      FL.Provider
+      FL.P2P
 
     -- Modules included in this library but not exported.
     -- other-modules:
@@ -71,10 +92,34 @@ library
     -- other-extensions:
 
     -- Other library packages from which modules are imported.
-    build-depends:    base ^>=4.18.3.0, stm, uuid, bytestring ^>=0.12.0.0, containers ^>=0.8, mtl ^>=2.3.0, binary ^>=0.8.9
+    build-depends:
+        base ^>=4.18.3.0,
+        stm,
+        uuid,
+        bytestring ^>=0.12.0.0,
+        containers ^>=0.8,
+        mtl ^>=2.3.0,
+        binary ^>=0.8.9,
+        network ^>=3.2,
+        hmatrix,
+        cassava,
+        vector,
+        random,
+        text,
+        time
 
     -- Directories containing source files.
     hs-source-dirs:   src
 
     -- Base language which the package is written in.
     default-language: GHC2021
+
+executable fl-actors
+    import:           warnings
+    main-is:          Main.hs
+    hs-source-dirs:   app
+    build-depends:
+        base ^>=4.18.3.0,
+        actors,
+        optparse-applicative
+    default-language: GHC2021

app/Main.hs

@@ -0,0 +1,150 @@
+module Main where
+
+import Control.Actor.Types (NodeAddr (..))
+import FL.P2P (runPeer)
+import FL.Provider (runCoordinator, runTrainer)
+import FL.Types
+import Options.Applicative
+
+-- ---------------------------------------------------------------------------
+-- CLI
+-- ---------------------------------------------------------------------------
+
+data Mode = Provider | P2P deriving (Show, Eq)
+
+data Role = Coordinator | Trainer | Peer deriving (Show, Eq)
+
+data CLIArgs = CLIArgs
+  { argMode         :: Mode
+  , argRole         :: Role
+  , argPort         :: Int
+  , argBasePort     :: Int
+  , argRounds       :: Int
+  , argEpochs       :: Int
+  , argLR           :: Double
+  , argBatch        :: Int
+  , argNumTrainers  :: Int
+  , argTrainerId    :: Int
+  , argCoordHost    :: String
+  , argCoordPort    :: Int
+  , argKnownPeers   :: [String]
+  , argDataPath     :: Maybe FilePath
+  , argPartition    :: PartitionMode
+  , argNorm         :: NormMode
+  }
+
+parseMode :: String -> Either String Mode
+parseMode "provider" = Right Provider
+parseMode "p2p"      = Right P2P
+parseMode s          = Left ("unknown mode: " <> s)
+
+parsePartition :: String -> Either String PartitionMode
+parsePartition "iid"    = Right IID
+parsePartition "noniid" = Right NonIID
+parsePartition s        = Left ("unknown partition: " <> s)
+
+parseNorm :: String -> Either String NormMode
+parseNorm "minmax" = Right MinMax
+parseNorm "zscore" = Right ZScore
+parseNorm s        = Left ("unknown norm: " <> s)
+
+parseRole :: String -> Either String Role
+parseRole "coordinator" = Right Coordinator
+parseRole "trainer"     = Right Trainer
+parseRole "peer"        = Right Peer
+parseRole s             = Left ("unknown role: " <> s)
+
+parsePeer :: String -> NodeAddr
+parsePeer s =
+  let (host, rest) = break (== ':') s
+  in NodeAddr host (read (drop 1 rest))
+
+argsParser :: Parser CLIArgs
+argsParser = CLIArgs
+  <$> option (eitherReader parseMode)
+        (long "mode" <> short 'm' <> metavar "provider|p2p"
+         <> help "Topology mode" <> value Provider <> showDefaultWith show)
+  <*> option (eitherReader parseRole)
+        (long "role" <> short 'r' <> metavar "coordinator|trainer|peer"
+         <> help "Node role" <> value Coordinator <> showDefaultWith show)
+  <*> option auto
+        (long "port" <> metavar "PORT"
+         <> help "This node's port" <> value 9100 <> showDefault)
+  <*> option auto
+        (long "base-port" <> metavar "PORT"
+         <> help "Base port (coordinator uses this)" <> value 9100 <> showDefault)
+  <*> option auto
+        (long "rounds" <> metavar "N"
+         <> help "Number of federated rounds" <> value 10 <> showDefault)
+  <*> option auto
+        (long "epochs" <> short 'e' <> metavar "N"
+         <> help "Local training epochs per round" <> value 5 <> showDefault)
+  <*> option auto
+        (long "lr" <> metavar "LR"
+         <> help "Learning rate" <> value 0.01 <> showDefault)
+  <*> option auto
+        (long "batch" <> metavar "N"
+         <> help "Mini-batch size" <> value 32 <> showDefault)
+  <*> option auto
+        (long "trainers" <> short 't' <> metavar "N"
+         <> help "Expected number of trainers/peers" <> value 3 <> showDefault)
+  <*> option auto
+        (long "id" <> metavar "N"
+         <> help "Trainer/peer index (0-based)" <> value 0 <> showDefault)
+  <*> strOption
+        (long "coord-host" <> metavar "HOST"
+         <> help "Coordinator hostname (trainer mode)" <> value "localhost")
+  <*> option auto
+        (long "coord-port" <> metavar "PORT"
+         <> help "Coordinator port (trainer mode)" <> value 9100 <> showDefault)
+  <*> many (strOption
+        (long "peer" <> metavar "HOST:PORT"
+         <> help "Known peer address (p2p mode, repeat for each peer)"))
+  <*> optional (strOption
+        (long "data" <> short 'd' <> metavar "PATH"
+         <> help "Path to PaySim CSV (uses synthetic data if absent)"))
+  <*> option (eitherReader parsePartition)
+        (long "partition" <> metavar "iid|noniid"
+         <> help "Data partition mode" <> value IID <> showDefaultWith show)
+  <*> option (eitherReader parseNorm)
+        (long "norm" <> metavar "minmax|zscore"
+         <> help "Feature normalization mode" <> value MinMax <> showDefaultWith show)
+
+main :: IO ()
+main = do
+  args <- execParser $
+    info (argsParser <**> helper)
+      ( fullDesc
+        <> progDesc "Federated Learning on the Haskell Actor Framework"
+        <> header "fl-actors — FedAvg fraud detection"
+      )
+
+  let cfg = FLConfig
+        { cfgRounds       = argRounds args
+        , cfgEpochs       = argEpochs args
+        , cfgLearningRate = argLR args
+        , cfgBatchSize    = argBatch args
+        , cfgNumTrainers  = argNumTrainers args
+        , cfgDataPath     = argDataPath args
+        , cfgPartition    = argPartition args
+        , cfgNorm         = argNorm args
+        , cfgBasePort     = argBasePort args
+        }
+
+  case (argMode args, argRole args) of
+    (Provider, Coordinator) ->
+      runCoordinator cfg { cfgBasePort = argPort args }
+
+    (Provider, Trainer) -> do
+      let coordAddr = NodeAddr (argCoordHost args) (fromIntegral (argCoordPort args))
+      runTrainer cfg (argTrainerId args) coordAddr
+
+    (P2P, Peer) -> do
+      let peers = map parsePeer (argKnownPeers args)
+          cfg'  = cfg { cfgBasePort = argPort args }
+      runPeer cfg' (argTrainerId args) peers
+
+    (mode, role) ->
+      putStrLn $ "Invalid combination: mode=" <> show mode <> " role=" <> show role
+        <> "\n  provider mode uses: --role=coordinator or --role=trainer"
+        <> "\n  p2p mode uses: --role=peer"

scripts/demo-p2p.sh

@@ -0,0 +1,85 @@
+#!/usr/bin/env bash
+#
+# Demo: federated learning in P2P mode (full mesh, no central coordinator).
+# Starts N peer nodes as separate processes; each trains locally, gossips
+# updates via PeerSync (CRDT-gated rounds), and aggregates independently.
+# At the end every peer must hold the SAME model — the script verifies the
+# saved model files are byte-identical.
+#
+# Usage:
+#   scripts/demo-p2p.sh [ROUNDS] [PEERS] [EPOCHS] [/path/to/paysim.csv]
+#
+# Without a CSV path, nodes fall back to synthetic data (fast; the metrics are
+# meaningless, but the gossip protocol and aggregation are fully exercised).
+set -euo pipefail
+
+ROUNDS="${1:-3}"
+PEERS="${2:-3}"
+EPOCHS="${3:-2}"
+DATA="${4:-}"
+
+BASE_PORT=9600
+OUT="demo-out/p2p"
+
+cd "$(dirname "$0")/.."
+echo "==> Building..."
+cabal build exe:fl-actors >/dev/null
+BIN="$(cabal list-bin fl-actors)"
+
+mkdir -p "$OUT"
+rm -f "$OUT"/*.log fl_model_peer-*.bin
+
+DATA_ARGS=()
+[ -n "$DATA" ] && DATA_ARGS=(--data "$DATA")
+
+PIDS=()
+cleanup() { kill "${PIDS[@]}" 2>/dev/null || true; wait 2>/dev/null || true; }
+trap cleanup EXIT INT TERM
+
+for i in $(seq 0 $((PEERS - 1))); do
+  PORT=$((BASE_PORT + i))
+  ARGS=(--mode p2p --role peer --id "$i" --port "$PORT"
+        --rounds "$ROUNDS" --trainers "$PEERS" --epochs "$EPOCHS")
+  for j in $(seq 0 $((PEERS - 1))); do
+    [ "$j" = "$i" ] && continue
+    ARGS+=(--peer "localhost:$((BASE_PORT + j))")
+  done
+  echo "==> Starting peer-$i on port $PORT"
+  "$BIN" "${ARGS[@]}" "${DATA_ARGS[@]}" > "$OUT/peer$i.log" 2>&1 &
+  PIDS+=($!)
+done
+
+echo "==> Federation running; waiting for all $PEERS peers to finish..."
+echo "    (live log: tail -f $OUT/peer0.log)"
+FAIL=0
+for pid in "${PIDS[@]}"; do
+  wait "$pid" || FAIL=1
+done
+PIDS=()
+if [ "$FAIL" -ne 0 ]; then
+  echo "!!  A peer exited abnormally; check $OUT/ logs"
+  exit 1
+fi
+
+echo
+echo "==> Evaluation results (peer-0's view):"
+grep "Eval round=" "$OUT/peer0.log" || true
+echo
+
+MODELS=(fl_model_peer-*.bin)
+if [ "${#MODELS[@]}" -ne "$PEERS" ]; then
+  echo "!!  Expected $PEERS model files, found ${#MODELS[@]}; check $OUT/ logs"
+  exit 1
+fi
+
+FIRST="${MODELS[0]}"
+for m in "${MODELS[@]:1}"; do
+  if ! cmp -s "$FIRST" "$m"; then
+    echo "!!  Divergence: $m differs from $FIRST"
+    md5sum "${MODELS[@]}"
+    exit 1
+  fi
+done
+echo "==> Success: all $PEERS peers converged to a byte-identical model:"
+ls -la "${MODELS[@]}"
+echo "    Full logs in $OUT/"

scripts/demo-provider.sh

@@ -0,0 +1,72 @@
+#!/usr/bin/env bash
+#
+# Demo: federated learning in PROVIDER mode (star topology).
+# Starts one coordinator and N trainer nodes as separate processes, runs the
+# federation, then prints the evaluation log and the saved model file.
+#
+# Usage:
+#   scripts/demo-provider.sh [ROUNDS] [TRAINERS] [EPOCHS] [/path/to/paysim.csv]
+#
+# Without a CSV path, nodes fall back to synthetic data (fast; the metrics are
+# meaningless, but the message flow and aggregation are fully exercised).
+set -euo pipefail
+
+ROUNDS="${1:-3}"
+TRAINERS="${2:-3}"
+EPOCHS="${3:-2}"
+DATA="${4:-}"
+
+BASE_PORT=9500
+OUT="demo-out/provider"
+
+cd "$(dirname "$0")/.."
+echo "==> Building..."
+cabal build exe:fl-actors >/dev/null
+BIN="$(cabal list-bin fl-actors)"
+
+mkdir -p "$OUT"
+rm -f "$OUT"/*.log fl_model_coordinator.bin
+
+DATA_ARGS=()
+[ -n "$DATA" ] && DATA_ARGS=(--data "$DATA")
+
+PIDS=()
+cleanup() { kill "${PIDS[@]}" 2>/dev/null || true; wait 2>/dev/null || true; }
+trap cleanup EXIT INT TERM
+
+echo "==> Starting coordinator on port $BASE_PORT (rounds=$ROUNDS, trainers=$TRAINERS, epochs=$EPOCHS)"
+"$BIN" --mode provider --role coordinator \
+  --port "$BASE_PORT" --rounds "$ROUNDS" --trainers "$TRAINERS" --epochs "$EPOCHS" \
+  "${DATA_ARGS[@]}" > "$OUT/coordinator.log" 2>&1 &
+COORD=$!
+PIDS+=("$COORD")
+sleep 1
+
+for i in $(seq 0 $((TRAINERS - 1))); do
+  echo "==> Starting trainer-$i on port $((BASE_PORT + i + 1))"
+  "$BIN" --mode provider --role trainer --id "$i" \
+    --base-port "$BASE_PORT" --coord-host localhost --coord-port "$BASE_PORT" \
+    --rounds "$ROUNDS" --trainers "$TRAINERS" --epochs "$EPOCHS" \
+    "${DATA_ARGS[@]}" > "$OUT/trainer$i.log" 2>&1 &
+  PIDS+=($!)
+done
+
+echo "==> Federation running; waiting for the coordinator to finish..."
+echo "    (live log: tail -f $OUT/coordinator.log)"
+if ! wait "$COORD"; then
+  echo "!!  Coordinator exited abnormally; last log lines:"
+  tail -20 "$OUT/coordinator.log"
+  exit 1
+fi
+
+echo
+echo "==> Evaluation results:"
+grep "Eval round=" "$OUT/coordinator.log" || true
+echo
+if [ -f fl_model_coordinator.bin ]; then
+  echo "==> Success: model saved to fl_model_coordinator.bin ($(stat -c %s fl_model_coordinator.bin) bytes)"
+  echo "    Full logs in $OUT/"
+else
+  echo "!!  No model file was produced; check $OUT/ logs"
+  exit 1
+fi

src/Control/Actor.hs

@@ -1,463 +1,78 @@
-{-# LANGUAGE StrictData #-}
+module Control.Actor
+  ( module Control.Actor.Types
+  , module Control.Actor.Transport
+  , module Control.Actor.Runtime
+  , module Control.Actor.Core
+  , module Control.Actor.Supervision
+  , module Control.Actor.Network
+  -- Initialization
+  , initRuntime
+  -- Demo
+  , pingActor
+  , forwardActorWithCell
+  , repeatActor
+  , system
+  , networkDemo
+  ) where
 
-module Control.Actor where
+import Control.Actor.Core
+import Control.Actor.Network
+import Control.Actor.Runtime
+import Control.Actor.Supervision
+import Control.Actor.Transport
+import Control.Actor.Types
 
-import Control.Concurrent (ThreadId, forkIO, threadDelay)
+import Control.Concurrent (forkIO, threadDelay)
-import Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar, takeMVar)
 import Control.Concurrent.STM
-  ( TMVar,
+  ( TMVar, TVar
-    TQueue,
+  , atomically, newEmptyTMVarIO, newTVarIO, readTMVar, readTVarIO, writeTQueue, writeTVar
-    TVar,
-    atomically,
-    flushTQueue,
-    modifyTVar,
-    newEmptyTMVarIO,
-    newTQueueIO,
-    newTVar,
-    newTVarIO,
-    orElse,
-    putTMVar,
-    readTMVar,
-    readTQueue,
-    readTVar,
-    readTVarIO,
-    tryTakeTMVar,
-    writeTQueue,
-    writeTVar,
   )
-import Control.Exception (AsyncException (..), SomeException, fromException, throwIO, try)
+import Control.Monad ((<=<), void)
-import Control.Monad (forM_, forever, void)
+import Control.Monad.Reader (MonadIO (..))
-import Control.Monad.Reader (MonadIO (liftIO), MonadReader (ask), ReaderT (runReaderT), asks, lift, withReaderT)
-import Data.Binary (Binary, decode, encode)
-import Data.ByteString.Lazy (ByteString)
-import Data.List (find)
-import Data.Map qualified as Map
-import Data.UUID (UUID)
-import Data.UUID.V4 (nextRandom)
-import GHC.Generics (Generic)
 import Unsafe.Coerce (unsafeCoerce)
+import Control.Exception (SomeException, try)
+import Control.Actor.Registry (createRegistry, registerActor, lookupRemoteActor)
 
-data NodeAddr = NodeAddr
+initRuntime :: NodeAddr -> (NodeAddr -> IO (Transport, NodeAddr)) -> IO Runtime
-  { nodeHost :: String,
+initRuntime myAddr createTransport = do
-    nodePort :: Integer
+  (transport, actualAddr) <- createTransport myAddr
-  }
+  rt0 <- newRuntime actualAddr transport
-  deriving (Eq, Show, Generic)
+  let rt = rt0 { rtSendRemote = \addr nm -> getOrCreateConn addr >>= cast' nm }
-
+  withRuntime rt $ void $ spawnActor sysHandlerFn stopOnDeath ()
-instance Binary NodeAddr
+  withRuntime rt $ void createRegistry
-
+  tListen transport $ \ch -> do
-type NodeId = Integer
+    result <- try @SomeException $ withRuntime rt $ handleNewConn ch
-
-thisNodeId :: NodeId
-thisNodeId = 0
-
-data ActorId = ActorId NodeId UUID
-  deriving (Eq, Ord, Show, Generic)
-
-instance Binary ActorId
-
-data ExitReason
-  = Normal
-  | Killed
-  | Exception SomeException
-  deriving (Show, Generic)
-
-data DeathMessage = DeathMessage
-  { dmActorId :: ActorId,
-    dmReason :: ExitReason
-  }
-  deriving (Show, Generic)
-
-data DeathTarget
-  = LocalTarget (TQueue DeathMessage)
-  | RemoteTarget ActorId
-
-data ActorState u = ActorState
-  { asId :: ActorId,
-    asLinks :: TVar [DeathTarget],
-    asEnv :: u
-  }
-
-newtype ActorM u r = ActorM
-  {unActorM :: ReaderT (ActorState u, Runtime) IO r}
-  deriving
-    ( Functor,
-      Applicative,
-      Monad,
-      MonadIO,
-      MonadReader (ActorState u, Runtime)
-    )
-
-runActorM :: ActorM u r -> Runtime -> ActorState u -> IO r
-runActorM m r s = runReaderT (unActorM m) (s, r)
-
-state :: ActorM u u
-state = asks (asEnv . fst)
-
-getSelf :: ActorM u SomeActorRef
-getSelf = do
-  (as, rt) <- ask
-  actors <- liftIO $ readTVarIO (rtActors rt)
-  let actorId = asId as
-      maybeRef = snd <$> Map.lookup actorId actors
-  case maybeRef of
-    Just ref -> return ref
-    Nothing  -> error "getSelf: actor not found in runtime"
-
-data ActorRef msg reply
-  = forall u. LocalRef
-      { arMsgQ :: TQueue (Envelope msg reply),
-        arDeathQ :: TQueue DeathMessage,
-        arState :: ActorState u
-      }
-  | RemoteRef ActorId
-
-data SomeActorRef = forall msg reply. SomeActorRef (ActorRef msg reply)
-
-someActorId :: SomeActorRef -> ActorId
-someActorId (SomeActorRef (LocalRef {arState})) = asId arState
-someActorId (SomeActorRef (RemoteRef aid)) = aid
-
-data SupervisorAction u
-  = Stop
-  | Continue u
-
-type CorrelationId = Integer
-
-data Envelope msg reply
-  = Cast msg
-  | Call msg (MVar (Maybe reply))
-
-data Runtime = Runtime
-  { rtNodeId :: NodeAddr,
-    rtActors :: TVar (Map.Map ActorId (ThreadId, SomeActorRef)),
-    rtPending :: TVar (Map.Map CorrelationId (MVar ByteString)),
-    rtNextCorr :: TVar CorrelationId,
-    rtNodeTable :: TVar (Map.Map NodeId NodeAddr),
-    rtTransport :: Transport
-  }
-
-data Transport = Transport
-  { sendBytes :: NodeAddr -> ByteString -> IO ()
-  }
-
-data RemoteEnvelope
-  = RemoteCast UUID ByteString
-  | RemoteCall UUID CorrelationId NodeAddr ByteString
-  | RemoteReply CorrelationId ByteString
-  deriving (Generic)
-
-instance Binary RemoteEnvelope
-
-newtype RuntimeM a = RuntimeM
-  {unRuntimeM :: ReaderT Runtime IO a}
-  deriving
-    ( Functor,
-      Applicative,
-      Monad,
-      MonadIO,
-      MonadReader Runtime
-    )
-
-liftRuntime :: RuntimeM a -> ActorM u a
-liftRuntime = ActorM . withReaderT snd . unRuntimeM
-
-withRuntime :: Runtime -> RuntimeM a -> IO a
-withRuntime = (. unRuntimeM) . flip runReaderT
-{-# INLINE withRuntime #-}
-
-newRuntime :: IO Runtime
-newRuntime = atomically $ do
-  actors <- newTVar Map.empty
-  pending <- newTVar Map.empty
-  nextCorr <- newTVar (0 :: Integer)
-  nodeTable <- newTVar Map.empty
-  return
-    Runtime
-      { rtNodeId = NodeAddr "localhost" 0,
-        rtActors = actors,
-        rtPending = pending,
-        rtNextCorr = nextCorr,
-        rtNodeTable = nodeTable,
-        rtTransport = Transport (\_ _ -> return ())
-      }
-
-lookupNode :: NodeId -> RuntimeM (Maybe NodeAddr)
-lookupNode nodeId = do
-  rt <- ask
-  liftIO $ atomically $ do
-    table <- readTVar $ rtNodeTable rt
-    return $ Map.lookup nodeId table
-
-cast' :: (Binary msg) => msg -> ActorRef msg reply -> RuntimeM ()
-cast' msg (LocalRef {arMsgQ}) = RuntimeM $ lift $ atomically $ writeTQueue arMsgQ (Cast msg)
-cast' msg (RemoteRef (ActorId nodeId uuid)) = do
-  rt <- ask
-  maybeAddr <- lookupNode nodeId
-  let payload = encode (RemoteCast uuid (encode msg))
-  case maybeAddr of
-    Just addr -> liftIO $ sendBytes (rtTransport rt) addr payload
-    Nothing -> liftIO $ putStrLn $ "cast: no node in lookup table with id " <> show nodeId
-
-cast :: (Binary msg) => msg -> ActorRef msg reply -> ActorM u ()
-cast = (liftRuntime .) . cast'
-
-castIn :: (Binary msg) => Int -> msg -> ActorRef msg reply -> ActorM u ()
-castIn ms msg ref = do
-  rt <- asks snd
-  liftIO $ void $ forkIO $ do
-    threadDelay (ms * 1000)
-    withRuntime rt $ cast' msg ref
-
-call' :: (Binary msg, Binary reply) => msg -> ActorRef msg reply -> RuntimeM (Maybe reply)
-call' msg (LocalRef {arMsgQ}) = liftIO $ do
-  mv <- newEmptyMVar
-  atomically $ writeTQueue arMsgQ (Call msg mv)
-  takeMVar mv
-call' msg (RemoteRef (ActorId nodeId uuid)) = do
-  rt <- ask
-  corrId <- liftIO $ atomically $ do
-    cid <- readTVar (rtNextCorr rt)
-    writeTVar (rtNextCorr rt) (cid + 1)
-    return cid
-  replyVar <- liftIO newEmptyMVar
-  liftIO $ atomically $ modifyTVar (rtPending rt) (Map.insert corrId replyVar)
-  let payload = encode (RemoteCall uuid corrId (rtNodeId rt) (encode msg))
-  maybeAddr <- lookupNode nodeId
-  case maybeAddr of
-    Just addr -> do
-      liftIO $ sendBytes (rtTransport rt) addr payload
-      raw <- liftIO $ takeMVar replyVar
-      liftIO $ atomically $ modifyTVar (rtPending rt) (Map.delete corrId)
-      return $ decode raw
-    Nothing -> do
-      liftIO $ putStrLn $ "call: no node in lookup table with id " <> show nodeId
-      return Nothing
-
-call :: (Binary msg, Binary reply) => msg -> ActorRef msg reply -> ActorM u (Maybe reply)
-call = (liftRuntime .) . call'
-
-type Actor u r = ActorM u (Maybe r, u)
-
-notifyOfDeath :: DeathMessage -> DeathTarget -> IO ()
-notifyOfDeath dm (LocalTarget q) = atomically $ writeTQueue q dm
-notifyOfDeath _ (RemoteTarget _) = return ()
-
-spawnActor ::
-  (m -> Actor u r) ->
-  (DeathMessage -> ActorM u (SupervisorAction u)) ->
-  u ->
-  RuntimeM (ActorRef m r)
-spawnActor actorFn deathFn initState = do
-  rt <- ask
-  mailbox <- liftIO newTQueueIO
-  deathQ <- liftIO newTQueueIO
-  links <- liftIO $ newTVarIO []
-  uuid <- liftIO nextRandom
-  let actorId = ActorId thisNodeId uuid
-      actorState = ActorState actorId links initState
-      actorRef = LocalRef mailbox deathQ actorState
-
-  let loop as = do
-        event <-
-          atomically $
-            (Left <$> readTQueue mailbox)
-              `orElse` (Right <$> readTQueue deathQ)
-        case event of
-          Left envelope ->
-            case envelope of
-              Cast msg -> do
-                (_, u') <- runActorM (actorFn msg) rt as
-                loop as {asEnv = u'}
-              Call msg mv -> do
-                (reply, u') <- runActorM (actorFn msg) rt as
-                putMVar mv reply
-                loop as {asEnv = u'}
-          Right dm -> do
-            action <- runActorM (deathFn dm) rt as
-            case action of
-              Stop -> return ()
-              Continue u -> loop as {asEnv = u}
-
-  tid <- liftIO $ forkIO $ do
-    result <- try (loop actorState) :: IO (Either SomeException ())
-    let reason = case result of
-          Right () -> Normal
-          Left exc -> case fromException exc of
-            Just ThreadKilled -> Killed
-            _anyOtherExc      -> Exception exc
-    links' <- readTVarIO (asLinks actorState)
-    let dm = DeathMessage actorId reason
-    forM_ links' (notifyOfDeath dm)
-    atomically $ modifyTVar (rtActors rt) (Map.delete actorId)
     case result of
-      Left exc -> throwIO exc
+      Left _   -> chClose ch
       Right () -> return ()
+  return rt
 
-  liftIO $ atomically $ modifyTVar (rtActors rt) (Map.insert actorId (tid, SomeActorRef actorRef))
-  return actorRef
 
-linkActorTo :: DeathTarget -> ActorRef m r -> RuntimeM ()
+-- Demo
-linkActorTo target (LocalRef {arState}) =
-  liftIO $ atomically $ modifyTVar (asLinks arState) (target :)
-linkActorTo _ (RemoteRef _) = return ()
 
-linkTo :: DeathTarget -> ActorM u ()
+pingActor :: Handler String () String
-linkTo target = do
+pingActor msg = continue ("Hello, " <> msg <> "!")
-  as <- asks fst
-  liftIO $ atomically $ modifyTVar (asLinks as) (target :)
 
-killActor :: ActorRef m r -> ActorM u ()
+forwardActorWithCell :: TMVar (ActorRef String String) -> Handler String () String
-killActor (LocalRef {arDeathQ, arState}) =
-  liftIO $ atomically $ writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
-killActor (RemoteRef _) =
-  liftIO $ putStrLn "killActor: remote kill not yet implemented"
-
-stopOnDeath :: DeathMessage -> ActorM u (SupervisorAction u)
-stopOnDeath _ = return Stop
-
--- Supervision
-
-data ChildSpec = forall m r. ChildSpec
-  { csRun :: DeathTarget -> RuntimeM (ActorRef m r),
-    csOnSpawn :: ActorRef m r -> IO ()
-  }
-
-child ::
-  (m -> Actor u r) ->
-  (DeathMessage -> ActorM u (SupervisorAction u)) ->
-  u ->
-  ChildSpec
-child msgFn deathFn initState =
-  ChildSpec
-    { csRun = \target -> do
-        ref <- spawnActor msgFn deathFn initState
-        linkActorTo target ref
-        return ref,
-      csOnSpawn = \_ -> return ()
-    }
-
-childWithRef ::
-  (m -> Actor u r) ->
-  (DeathMessage -> ActorM u (SupervisorAction u)) ->
-  u ->
-  TMVar (ActorRef m r) ->
-  ChildSpec
-childWithRef msgFn deathFn initState cell =
-  ChildSpec
-    { csRun = \target -> do
-        ref <- spawnActor msgFn deathFn initState
-        linkActorTo target ref
-        return ref,
-      csOnSpawn = \ref -> atomically $ do
-        void $ tryTakeTMVar cell
-        putTMVar cell ref
-    }
-
-data RestartStrategy = OneForOne | OneForAll | RestForOne
-
-data ChildSlot = forall m r. ChildSlot
-  { slotSpec :: ChildSpec,
-    slotRef :: ActorRef m r,
-    slotId :: ActorId
-  }
-
-spawnSlot :: DeathTarget -> ChildSpec -> RuntimeM ChildSlot
-spawnSlot target spec@ChildSpec{csRun, csOnSpawn} = do
-  ref <- csRun target
-  _   <- liftIO $ csOnSpawn ref
-  return $ ChildSlot spec ref (someActorId (SomeActorRef ref))
-
-supervise' :: RestartStrategy -> [ChildSpec] -> RuntimeM ()
-supervise' strategy specs = do
-  rt <- ask
-  supDeathQ <- liftIO newTQueueIO
-  let target = LocalTarget supDeathQ
-  slots <- mapM (spawnSlot target) specs
-  slotsVar <- liftIO $ newTVarIO slots
-  _ <- liftIO $ forkIO $ forever $ do
-    DeathMessage deadId _ <- atomically $ readTQueue supDeathQ
-    slots' <- readTVarIO slotsVar
-    case strategy of
-      OneForOne -> doOneForOne rt target slotsVar slots' deadId
-      OneForAll -> doOneForAll rt target slotsVar supDeathQ slots'
-      RestForOne -> doRestForOne rt target slotsVar supDeathQ slots' deadId
-  return ()
-
-supervise :: RestartStrategy -> [ChildSpec] -> ActorM u ()
-supervise = (liftRuntime .) . supervise'
-
-doOneForOne ::
-  Runtime -> DeathTarget -> TVar [ChildSlot] -> [ChildSlot] -> ActorId -> IO ()
-doOneForOne rt target slotsVar slots deadId =
-  case find (\s -> slotId s == deadId) slots of
-    Nothing -> return ()
-    Just slot -> do
-      newSlot <- withRuntime rt $ spawnSlot target (slotSpec slot)
-      atomically $
-        modifyTVar slotsVar $
-          map (\s -> if slotId s == deadId then newSlot else s)
-
-doOneForAll ::
-  Runtime ->
-  DeathTarget ->
-  TVar [ChildSlot] ->
-  TQueue DeathMessage ->
-  [ChildSlot] ->
-  IO ()
-doOneForAll rt target slotsVar supDeathQ slots = do
-  mapM_ (killSlot supDeathQ) slots
-  atomically $ void $ flushTQueue supDeathQ
-  newSlots <- withRuntime rt $ mapM (spawnSlot target . slotSpec) slots
-  atomically $ writeTVar slotsVar newSlots
-
-doRestForOne ::
-  Runtime ->
-  DeathTarget ->
-  TVar [ChildSlot] ->
-  TQueue DeathMessage ->
-  [ChildSlot] ->
-  ActorId ->
-  IO ()
-doRestForOne rt target slotsVar supDeathQ slots deadId = do
-  let (before, fromDead) = break (\s -> slotId s == deadId) slots
-  case fromDead of
-    []        -> return ()
-    _nonempty -> do
-      mapM_ (killSlot supDeathQ) (drop 1 fromDead)
-      atomically $ void $ flushTQueue supDeathQ
-      newSlots <- withRuntime rt $ mapM (spawnSlot target . slotSpec) fromDead
-      atomically $ writeTVar slotsVar (before ++ newSlots)
-
-killSlot :: TQueue DeathMessage -> ChildSlot -> IO ()
-killSlot _ (ChildSlot {slotRef, slotId}) = case slotRef of
-  LocalRef {arDeathQ} -> atomically $ writeTQueue arDeathQ (DeathMessage slotId Killed)
-  RemoteRef _ -> return ()
-
------ Demo
-
-pingActor :: String -> Actor () String
-pingActor msg = return (Just ("Hello, " <> msg <> "!"), ())
-
-forwardActorWithCell :: TMVar (ActorRef String String) -> String -> Actor () String
 forwardActorWithCell cell msg = do
   pingRef <- liftIO $ atomically $ readTMVar cell
   reply   <- call msg pingRef
   case reply of
     Nothing -> liftIO $ putStrLn "forwardActorWithCell: received empty reply!"
     Just x  -> liftIO $ putStrLn ("forwardActorWithCell: received reply - " <> x)
-  return (reply, ())
+  maybeContinue reply
 
-repeatActor :: String -> TMVar (ActorRef String String) -> () -> Actor () String
+repeatActor :: String -> TMVar (ActorRef String String) -> Handler () () String
 repeatActor r cell () = do
   ref             <- liftIO $ atomically $ readTMVar cell
   (SomeActorRef self) <- getSelf
   cast r ref
   castIn 1000 () (unsafeCoerce self)
-  return (Nothing, ())
+  pass
 
 system :: IO ()
 system = do
-  rt <- newRuntime
+  rt <- initRuntime (NodeAddr "localhost" 9000) createTCPTransport
   withRuntime rt $ do
     pingCell    <- liftIO newEmptyTMVarIO
     forwardCell <- liftIO newEmptyTMVarIO
@@ -465,12 +80,127 @@ system = do
 
     _ <- liftIO $ forkIO $ do
       repeatRef <- atomically $ readTMVar repeatCell
-      withRuntime rt $ do
+      withRuntime rt $
         cast' () repeatRef
 
     supervise'
       OneForOne
-      [ childWithRef pingActor stopOnDeath () pingCell,
+      [ childWithRef pingActor stopOnDeath () pingCell
-        childWithRef (forwardActorWithCell pingCell) stopOnDeath () forwardCell,
+      , childWithRef (forwardActorWithCell pingCell) stopOnDeath () forwardCell
-        childWithRef (repeatActor "repeaaat" forwardCell) stopOnDeath () repeatCell
+      , childWithRef (repeatActor "repeaaat" forwardCell) stopOnDeath () repeatCell
       ]
+
+-- Network demo
+
+newNodeActor :: Handler NodeAddr () NodeId
+newNodeActor = continue <=< liftRuntime . connect Nothing
+
+-- | Actor on node 2: echo — returns whatever it received.
+echoActor :: Handler String () String
+echoActor = continue
+
+-- | Actor on node 2: printer — side-effects only.
+printerActor :: Handler String () ()
+printerActor msg = liftIO (putStrLn $ "  [remote/print] " <> msg) >> pass
+
+-- | Trivial actor used as a stub when we only care about death notifications.
+noopActor :: Handler () () ()
+noopActor _ = pass
+
+-- | Captures the received string into a TVar (used as a node-1 receiver actor).
+recvActorFn :: TVar (Maybe String) -> Handler String () ()
+recvActorFn var msg =
+  liftIO (atomically $ writeTVar var (Just msg)) >> pass
+
+-- | End-to-end networking demo.
+--
+-- Covers:
+--   1. Remote cast  (node 1 → node 2)
+--   2. Remote call  (node 1 ↔ node 2, request/reply)
+--   3. Reverse cast (node 2 → node 1)
+--   4. Cross-node death notification
+networkDemo :: IO ()
+networkDemo = do
+  putStrLn "=== networking demo ==="
+
+  -- Port 0 lets the OS pick a free port each run, avoiding stale-socket
+  -- conflicts when re-running in the same GHCi session.
+  rt1 <- initRuntime (NodeAddr "localhost" 0) createTCPTransport
+  rt2 <- initRuntime (NodeAddr "localhost" 0) createTCPTransport
+  let addr1 = rtNodeId rt1
+      addr2 = rtNodeId rt2
+
+  -- connect returns the locally-assigned NodeId for the peer.
+  -- NodeId 0 is always self; remotes get ids ≥ 1.
+  node2 <- withRuntime rt1 $ connect Nothing addr2
+  node1 <- withRuntime rt2 $ connect Nothing addr1
+
+  -- Spawn and register workers on node 2.
+  echoRef  <- withRuntime rt2 $ spawnActor echoActor   stopOnDeath ()
+  printRef <- withRuntime rt2 $ spawnActor printerActor stopOnDeath ()
+  withRuntime rt2 $ registerActor "echo"    echoRef
+  withRuntime rt2 $ registerActor "printer" printRef
+
+  -- Small delay to let NMRegistryUpdate propagate to node 1.
+  threadDelay 100_000
+
+  -- 1. Remote cast: fire-and-forget from node 1 to node 2.
+  putStrLn "\n[1] remote cast  node 1 -> node 2"
+  Just remotePrint <- withRuntime rt1 $ lookupRemoteActor "printer"
+  withRuntime rt1 $ cast' "hello from node 1" remotePrint
+  threadDelay 300_000
+
+  -- 2. Remote call: request/reply across nodes.
+  putStrLn "\n[2] remote call  node 1 <-> node 2"
+  Just remoteEcho' <- withRuntime rt1 $ lookupRemoteActor "echo"
+  let
+    remoteEcho :: ActorRef String String
+    remoteEcho = (unsafeCoerce remoteEcho')
+  reply <- withRuntime rt1 $ call' "ping" remoteEcho
+  putStrLn $ "    reply: " <> show reply
+
+  -- 3. Reverse cast: actor on node 2 sends to an actor on node 1.
+  putStrLn "\n[3] reverse cast  node 2 -> node 1"
+  recvVar <- newTVarIO (Nothing :: Maybe String)
+  recvRef <- withRuntime rt1 $ spawnActor (recvActorFn recvVar) stopOnDeath ()
+  withRuntime rt1 $ registerActor "receiver" recvRef
+  threadDelay 100_000
+  Just remoteRecv <- withRuntime rt2 $ lookupRemoteActor "receiver"
+  withRuntime rt2 $ cast' "hello from node 2" remoteRecv
+  threadDelay 300_000
+  readTVarIO recvVar >>= \v -> putStrLn $ "    received: " <> show v
+
+  -- 4. Cross-node death notification.
+  --
+  -- The mortal actor lives on node 2.  When it dies, it sends NMDeath to
+  -- node 1.  The watcher actor on node 1 has registered interest in the
+  -- mortal's ActorId; routeRemoteDeath finds it and delivers the message.
+  putStrLn "\n[4] cross-node death notification"
+  diedVar  <- newTVarIO False
+  watchRef <- withRuntime rt1 $ spawnActor noopActor
+    (\dm -> do
+        liftIO $ putStrLn $ "  [watcher] death of " <> show (dmActorId dm)
+        liftIO $ atomically $ writeTVar diedVar True
+        return Stop)
+    ()
+  mortalRef <- withRuntime rt2 $ spawnActor noopActor stopOnDeath ()
+
+  let mortalId = someActorId (SomeActorRef mortalRef)
+      watchId  = someActorId (SomeActorRef watchRef)
+
+  -- mortalRef's links: on death, send NMDeath to node 1 (by NodeId).
+  withRuntime rt2 $ linkActorTo (RemoteTarget watchId  node1) mortalRef
+  -- watchRef's links: declare interest in mortalId so routeRemoteDeath
+  -- on node 1 can find this actor when NMDeath arrives.
+  withRuntime rt1 $ linkActorTo (RemoteTarget mortalId node2) watchRef
+
+  -- Kill the mortal actor by posting directly to its death queue.
+  case mortalRef of
+    LocalRef {arDeathQ, arState} ->
+      atomically $ writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
+    RemoteRef _ -> return ()
+
+  threadDelay 500_000
+  readTVarIO diedVar >>= \d -> putStrLn $ "    death delivered: " <> show d
+
+  putStrLn "\n=== done ==="

src/Control/Actor/Core.hs

@@ -0,0 +1,289 @@
+{-# LANGUAGE StrictData #-}
+
+module Control.Actor.Core
+  ( ActorM (..)
+  , runActorM
+  , state
+  , getSelf
+  , Actor
+  , ActorResult (..)
+  , Handler
+  , liftRuntime
+  , notifyOfDeath
+  , spawnActor
+  , spawnActorAs
+  , linkActorTo
+  , linkTo
+  , killActor
+  , stopOnDeath
+  , cast'
+  , call'
+  , cast
+  , call
+  , castIn
+  , pass
+  , passWith
+  , continue
+  , maybeContinue
+  , become
+  , lastMessageFrom
+  ) where
+
+import Control.Actor.Runtime
+import Control.Actor.Types
+import Control.Concurrent (forkIO, killThread, threadDelay)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Concurrent.STM
+  ( atomically
+  , modifyTVar
+  , newTQueueIO
+  , newTVarIO
+  , orElse
+  , readTQueue
+  , readTVar
+  , readTVarIO
+  , writeTQueue
+  , writeTVar
+  )
+import Control.Exception
+  ( AsyncException (..)
+  , SomeException
+  , fromException
+  , throwIO
+  , try
+  )
+import Control.Monad (forM_, join, void)
+import Data.Maybe (fromMaybe)
+import System.Timeout (timeout)
+import Control.Monad.Reader
+  ( MonadIO (..)
+  , MonadReader (..)
+  , ReaderT (..)
+  , asks
+  , withReaderT
+  )
+import Data.Binary (Binary, decode, encode)
+import Data.Map qualified as Map
+import Data.UUID (UUID)
+import Data.UUID.V4 (nextRandom)
+
+newtype ActorM u r = ActorM
+  { unActorM :: ReaderT (ActorState u, Runtime) IO r }
+  deriving
+    ( Functor, Applicative, Monad, MonadIO
+    , MonadReader (ActorState u, Runtime)
+    )
+
+runActorM :: ActorM u r -> Runtime -> ActorState u -> IO r
+runActorM m rt s = runReaderT (unActorM m) (s, rt)
+
+state :: ActorM u u
+state = asks (asEnv . fst)
+
+lastMessageFrom :: ActorM u NodeId
+lastMessageFrom = asks (asLatestFrom . fst)
+
+getSelf :: ActorM u SomeActorRef
+getSelf = do
+  (as, rt) <- ask
+  actors <- liftIO $ readTVarIO (rtActors rt)
+  let actorId  = asId as
+      maybeRef = snd <$> Map.lookup actorId actors
+  case maybeRef of
+    Just ref -> return ref
+    Nothing  -> error "getSelf: actor not found in runtime"
+
+data ActorResult msg u r = ActorResult
+  { actorReply  :: Maybe r
+  , actorState  :: u
+  , actorBecome :: Maybe (msg -> ActorM u (ActorResult msg u r))
+  }
+
+type Actor msg u r = ActorM u (ActorResult msg u r)
+
+type Handler msg u r = msg -> Actor msg u r
+
+liftRuntime :: RuntimeM a -> ActorM u a
+liftRuntime = ActorM . withReaderT snd . unRuntimeM
+
+notifyOfDeath :: Runtime -> DeathMessage -> DeathTarget -> IO ()
+notifyOfDeath _  dm (LocalTarget q)         = atomically $ writeTQueue q dm
+notifyOfDeath rt dm (RemoteTarget _ nodeId) =
+  withRuntime rt $ do
+    maybeAddr <- lookupNode nodeId
+    case maybeAddr of
+      Nothing   -> return ()
+      Just addr -> rtSendRemote rt addr (NMDeath (dmActorId dm) (toRemoteExitReason (dmReason dm)))
+
+spawnActorAs ::
+  (Binary m, Binary r) =>
+  UUID ->
+  Handler m u r ->
+  (DeathMessage -> ActorM u (SupervisorAction u)) ->
+  u ->
+  RuntimeM (ActorRef m r)
+spawnActorAs uuid actorFn deathFn initState = do
+  liftIO $ putStrLn $ "spawning actor with UUID " <> show uuid
+  rt      <- ask
+  mailbox <- liftIO newTQueueIO
+  deathQ  <- liftIO newTQueueIO
+  links   <- liftIO $ newTVarIO []
+  let actorId    = ActorId thisNodeId uuid
+      actorState = ActorState actorId links initState thisNodeId
+      actorRef   = LocalRef mailbox deathQ actorState
+
+  let loop fn as = do
+        event <-
+          atomically $
+            (Left <$> readTQueue mailbox)
+              `orElse` (Right <$> readTQueue deathQ)
+        case event of
+          Left envelope ->
+            case envelope of
+              (from, Cast msg) -> do
+                ActorResult _ u' mFn <- runActorM (fn msg) rt as {asLatestFrom = from}
+                loop (fromMaybe fn mFn) as {asEnv = u'}
+              (from, Call msg mv) -> do
+                ActorResult reply u' mFn <- runActorM (fn msg) rt as {asLatestFrom = from}
+                putMVar mv reply
+                loop (fromMaybe fn mFn) as {asEnv = u'}
+          Right dm -> do
+            action <- runActorM (deathFn dm) rt as
+            case action of
+              Stop       -> return ()
+              Continue u -> loop fn as {asEnv = u}
+
+  -- The thread blocks on `ready` until it is registered in rtActors:
+  -- otherwise messages dispatched by UUID in the window before the insert
+  -- would be dropped, and a death racing the insert would leave a stale
+  -- entry behind (the death path's delete running before the insert).
+  ready <- liftIO newEmptyMVar
+  tid <- liftIO $ forkIO $ do
+    takeMVar ready
+    result <- try @SomeException (loop actorFn actorState)
+    let reason = case result of
+          Right () -> Normal
+          Left exc -> case fromException exc of
+            Just ThreadKilled -> Killed
+            _anyOtherExc      -> Exception exc
+    links' <- readTVarIO (asLinks actorState)
+    forM_ links' (notifyOfDeath rt (DeathMessage actorId reason))
+    atomically $ modifyTVar (rtActors rt) (Map.delete actorId)
+    case result of
+      Left exc -> throwIO exc
+      Right () -> return ()
+
+  liftIO $ atomically $
+    modifyTVar (rtActors rt) (Map.insert actorId (tid, SomeActorRef actorRef))
+  liftIO $ putMVar ready ()
+  return actorRef
+
+spawnActor ::
+  (Binary m, Binary r) =>
+  Handler m u r ->
+  (DeathMessage -> ActorM u (SupervisorAction u)) ->
+  u ->
+  RuntimeM (ActorRef m r)
+spawnActor actorFn deathFn initState = do
+  uuid <- liftIO nextRandom
+  spawnActorAs uuid actorFn deathFn initState
+
+linkActorTo :: DeathTarget -> ActorRef m r -> RuntimeM ()
+linkActorTo target (LocalRef {arState}) =
+  liftIO $ atomically $ modifyTVar (asLinks arState) (target :)
+linkActorTo _ (RemoteRef _) = return ()
+
+linkTo :: DeathTarget -> ActorM u ()
+linkTo target = do
+  as <- asks fst
+  liftIO $ atomically $ modifyTVar (asLinks as) (target :)
+
+-- | Forcibly terminate an actor's thread. Previously this only posted a
+-- 'DeathMessage' to the target's death queue, which a deathFn returning
+-- 'Continue' would silently ignore — and even on 'Stop' the links were told
+-- the exit was 'Normal'. Killing the thread makes links see 'Killed'.
+killActor :: ActorRef m r -> ActorM u ()
+killActor ref@(LocalRef {}) = do
+  rt <- asks snd
+  liftIO $ do
+    actors <- readTVarIO (rtActors rt)
+    forM_ (Map.lookup (actorRefId ref) actors) (killThread . fst)
+killActor (RemoteRef _) =
+  liftIO $ putStrLn "killActor: remote kill not yet implemented"
+
+stopOnDeath :: DeathMessage -> ActorM u (SupervisorAction u)
+stopOnDeath _ = return Stop
+
+cast' :: (Binary msg) => msg -> ActorRef msg reply -> RuntimeM ()
+cast' msg (LocalRef {arMsgQ}) =
+  liftIO $ atomically $ writeTQueue arMsgQ (thisNodeId, Cast msg)
+cast' msg (RemoteRef (ActorId nodeId uuid)) = do
+  rt <- ask
+  maybeAddr <- lookupNode nodeId
+  case maybeAddr of
+    Nothing   -> liftIO $ putStrLn $ "cast: no node in lookup table with id " <> show nodeId
+    Just addr -> do
+      result <- liftIO $ try @SomeException $ withRuntime rt $ rtSendRemote rt addr (NMCast uuid (encode msg))
+      case result of
+        Left e  -> liftIO $ putStrLn $ "cast: send to " <> show addr <> " failed: " <> show e
+        Right _ -> return ()
+
+-- | How long a call waits for a reply before giving up with Nothing.
+-- Without this, a call to a dead/missing actor (or over a dead connection)
+-- would block the caller forever.
+callTimeoutMicros :: Int
+callTimeoutMicros = 10_000_000
+
+call' :: (Binary msg, Binary reply) => msg -> ActorRef msg reply -> RuntimeM (Maybe reply)
+call' msg (LocalRef {arMsgQ}) = liftIO $ do
+  mv <- newEmptyMVar
+  atomically $ writeTQueue arMsgQ (thisNodeId, Call msg mv)
+  join <$> timeout callTimeoutMicros (takeMVar mv)
+call' msg (RemoteRef (ActorId nodeId uuid)) = do
+  rt <- ask
+  corrId <- liftIO $ atomically $ do
+    cid <- readTVar (rtNextCorr rt)
+    writeTVar (rtNextCorr rt) (cid + 1)
+    return cid
+  replyVar <- liftIO newEmptyMVar
+  liftIO $ atomically $ modifyTVar (rtPending rt) (Map.insert corrId replyVar)
+  maybeAddr <- lookupNode nodeId
+  case maybeAddr of
+    Nothing -> liftIO $ do
+      putStrLn $ "call: no node in lookup table with id " <> show nodeId
+      atomically $ modifyTVar (rtPending rt) (Map.delete corrId)
+      return Nothing
+    Just addr -> do
+      rtSendRemote rt addr (NMCall uuid corrId (rtNodeId rt) (encode msg))
+      raw <- liftIO $ timeout callTimeoutMicros (takeMVar replyVar)
+      liftIO $ atomically $ modifyTVar (rtPending rt) (Map.delete corrId)
+      return $ decode <$> raw
+
+cast :: (Binary msg) => msg -> ActorRef msg reply -> ActorM u ()
+cast = (liftRuntime .) . cast'
+
+castIn :: (Binary msg) => Int -> msg -> ActorRef msg reply -> ActorM u ()
+castIn ms msg ref = do
+  rt <- asks snd
+  liftIO $ void $ forkIO $ do
+    threadDelay (ms * 1000)
+    withRuntime rt $ cast' msg ref
+
+call :: (Binary msg, Binary reply) => msg -> ActorRef msg reply -> ActorM u (Maybe reply)
+call = (liftRuntime .) . call'
+
+continue :: r -> Actor msg u r
+continue x = (\u -> ActorResult (Just x) u Nothing) <$> state
+
+maybeContinue :: Maybe r -> Actor msg u r
+maybeContinue Nothing = pass
+maybeContinue (Just r) = continue r
+
+pass :: Actor msg u r
+pass = (\u -> ActorResult Nothing u Nothing) <$> state
+
+passWith :: u -> Actor msg u r
+passWith u = return (ActorResult Nothing u Nothing)
+
+become :: Handler msg u r -> Actor msg u r
+become fn = (\u -> ActorResult Nothing u (Just fn)) <$> state

src/Control/Actor/Network.hs

@@ -0,0 +1,306 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE StrictData #-}
+
+module Control.Actor.Network
+  ( spawnConnTree
+  , handleNewConn
+  , getOrCreateConn
+  , connect
+  , findByUUID
+  , routeRemoteDeath
+  , validateAndDispatch
+  , sysHandlerFn
+  , createTCPTransport
+  ) where
+
+import Control.Actor.Core
+  ( ActorM, ActorResult (..), Handler, cast', liftRuntime, linkActorTo
+  , pass, spawnActor, state, stopOnDeath )
+import Control.Actor.Registry (deregisterNode)
+import Control.Actor.Runtime (Runtime (..), RuntimeM, addrToNodeId, withRuntime)
+import Control.Actor.Transport (ConnHandle (..), Transport (..), createTCPTransport)
+import Control.Actor.Types
+import Control.Concurrent (forkIO)
+import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
+import Control.Concurrent.STM
+  ( atomically
+  , modifyTVar
+  , newEmptyTMVar
+  , readTMVar
+  , readTVar
+  , readTVarIO
+  , tryPutTMVar
+  , writeTQueue
+  , writeTVar
+  )
+import Control.Exception (SomeException, throwIO, toException, try)
+import Control.Monad (forM_, forever, unless, void, when)
+import Control.Monad.Reader (MonadIO (..), MonadReader (..), asks)
+import Data.Binary (decode, decodeOrFail, encode)
+import Data.ByteString.Lazy (ByteString)
+import Data.Map qualified as Map
+import Data.Maybe (fromMaybe)
+import System.Timeout (timeout)
+
+-- Connection actors
+
+connActorFn :: Handler NetworkMessage ConnHandle ()
+connActorFn nm = do
+  ch <- state
+  liftIO $ chSend ch (encode nm)
+  return $ ActorResult Nothing ch Nothing
+
+connDeathFn :: NodeAddr -> DeathMessage -> ActorM ConnHandle (SupervisorAction ConnHandle)
+connDeathFn peer _ = do
+  ch  <- state
+  aid <- asks (asId . fst)
+  rt  <- liftRuntime ask
+  -- Only clear the pool entry if it is ours: with simultaneous connects two
+  -- conn actors can exist for the same peer, and the loser must not evict
+  -- the active one.
+  wasActive <- liftIO $ atomically $ do
+    conns <- readTVar (rtConnections rt)
+    case Map.lookup peer conns of
+      Just ref | actorRefId ref == aid -> do
+        modifyTVar (rtConnections rt) (Map.delete peer)
+        return True
+      _else -> return False
+  liftIO $ chClose ch
+  when wasActive $ do
+    mNodeId <- liftRuntime $ addrToNodeId peer
+    forM_ mNodeId $ \nodeId -> liftRuntime $ deregisterNode nodeId
+  return Stop
+
+routerActorFn :: NodeAddr -> Handler ByteString () ()
+routerActorFn senderAddr raw =
+  case decodeOrFail raw of
+    Left _ -> do
+      liftIO $ putStrLn $ "router: dropping undecodable frame from " <> show senderAddr
+      pass
+    Right (_, _, nm) -> do
+      valid <- liftRuntime $ validateAndDispatch senderAddr nm
+      unless valid $ liftIO $
+        putStrLn $ "router: dropping invalid message from " <> show senderAddr
+      pass
+
+-- Connection supervision tree
+
+spawnConnTree :: NodeAddr -> ConnHandle -> RuntimeM (ActorRef NetworkMessage ())
+spawnConnTree peer ch = do
+  rt        <- ask
+  routerRef <- spawnActor (routerActorFn peer) stopOnDeath ()
+  connRef   <- spawnActor connActorFn (connDeathFn peer) ch
+  -- When conn dies (peer disconnected), notify router so it shuts down too
+  case routerRef of
+    LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) connRef
+    RemoteRef _         -> return ()
+  -- And vice versa: if the router dies, close the connection down cleanly
+  case connRef of
+    LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) routerRef
+    RemoteRef _         -> return ()
+  liftIO $ case connRef of
+    LocalRef {arDeathQ, arState} ->
+      void $ forkIO $ do
+        result <- try @SomeException $ forever $
+          chRecv ch >>= \raw -> withRuntime rt $ cast' raw routerRef
+        case result of
+          Left _   -> atomically $
+            writeTQueue arDeathQ (DeathMessage (asId arState) Killed)
+          Right () -> return ()
+    RemoteRef _ -> return ()
+  return connRef
+
+-- Incoming connection handler
+
+handleNewConn :: ConnHandle -> RuntimeM ()
+handleNewConn ch = do
+  rt  <- ask
+  raw <- liftIO $ chRecv ch
+  case decode raw :: NetworkMessage of
+    NMHandshake peerAddr -> do
+      ref <- spawnConnTree peerAddr ch
+      liftIO $ atomically $ do
+        conns <- readTVar (rtConnections rt)
+        -- On a simultaneous connect two sockets (and two conn trees) exist
+        -- for the same peer. Only the pool entry is deduplicated; the extra
+        -- tree must stay alive because it is the read path for whatever the
+        -- peer sends on its socket. It cleans itself up when that socket
+        -- closes. Local waiters get the pool winner, never the extra ref.
+        let winner = Map.findWithDefault ref peerAddr conns
+        unless (Map.member peerAddr conns) $
+          modifyTVar (rtConnections rt) (Map.insert peerAddr ref)
+        table <- readTVar (rtNodeTable rt)
+        unless (any (== peerAddr) (Map.elems table)) $ do
+          nid <- readTVar (rtNextNodeId rt)
+          writeTVar (rtNextNodeId rt) (nid + 1)
+          modifyTVar (rtNodeTable rt) (Map.insert nid peerAddr)
+        promises <- readTVar (rtConnPromises rt)
+        case Map.lookup peerAddr promises of
+          Nothing -> return ()
+          Just p  -> do
+            modifyTVar (rtConnPromises rt) (Map.delete peerAddr)
+            void $ tryPutTMVar p (Right winner)
+    _else -> liftIO $ chClose ch
+
+-- Connection pool
+
+getOrCreateConn :: NodeAddr -> RuntimeM (ActorRef NetworkMessage ())
+getOrCreateConn peer = do
+  rt <- ask
+  action <- liftIO $ atomically $ do
+    conns <- readTVar (rtConnections rt)
+    case Map.lookup peer conns of
+      Just ref -> return (Left ref)
+      Nothing  -> do
+        promises <- readTVar (rtConnPromises rt)
+        case Map.lookup peer promises of
+          Just p  -> return (Right (Left p))
+          Nothing -> do
+            p <- newEmptyTMVar
+            modifyTVar (rtConnPromises rt) (Map.insert peer p)
+            return (Right (Right p))
+  case action of
+    Left ref              -> return ref
+    Right (Left promise)  -> do
+      -- Someone else is connecting; wait for their outcome so a failed
+      -- connect propagates instead of blocking us forever.
+      outcome <- liftIO $ atomically $ readTMVar promise
+      either (liftIO . throwIO) return outcome
+    Right (Right promise) -> do
+      result <- liftIO $ try @SomeException $ tConnect (rtTransport rt) peer
+      case result of
+        Left e -> do
+          liftIO $ atomically $ do
+            modifyTVar (rtConnPromises rt) (Map.delete peer)
+            void $ tryPutTMVar promise (Left e)
+          liftIO $ throwIO e
+        Right ch -> do
+          liftIO $ chSend ch (encode (NMHandshake (rtNodeId rt)))
+          ref <- spawnConnTree peer ch
+          liftIO $ atomically $ do
+            conns <- readTVar (rtConnections rt)
+            unless (Map.member peer conns) $
+              modifyTVar (rtConnections rt) (Map.insert peer ref)
+            modifyTVar (rtConnPromises rt) (Map.delete peer)
+            void $ tryPutTMVar promise (Right ref)
+          return ref
+
+-- Message dispatch
+
+routeRemoteDeath :: NodeAddr -> ActorId -> RemoteExitReason -> RuntimeM ()
+routeRemoteDeath senderAddr (ActorId _ uuid) reason = do
+  rt           <- ask
+  mSenderNid   <- addrToNodeId senderAddr
+  liftIO $ do
+    let deadId     = ActorId (fromMaybe 0 mSenderNid) uuid
+        exitReason = case reason of
+          RNormal      -> Normal
+          RKilled      -> Killed
+          RException s -> Exception (toException (userError s))
+        dm = DeathMessage deadId exitReason
+    actors <- readTVarIO (rtActors rt)
+    forM_ (Map.elems actors) $ \(_, SomeActorRef ref) ->
+      case ref of
+        LocalRef {arDeathQ, arState} -> do
+          links <- readTVarIO (asLinks arState)
+          forM_ links $ \case
+            RemoteTarget (ActorId _ uid) nodeId
+              | uid == uuid && Just nodeId == mSenderNid ->
+                  atomically $ writeTQueue arDeathQ dm
+            _else -> return ()
+        RemoteRef _ -> return ()
+
+validateAndDispatch :: NodeAddr -> NetworkMessage -> RuntimeM Bool
+validateAndDispatch senderAddr nm = case nm of
+
+  NMHandshake _ -> return False
+
+  NMReply corrId payload -> do
+    rt <- ask
+    liftIO $ do
+      pending <- readTVarIO (rtPending rt)
+      case Map.lookup corrId pending of
+        Nothing -> return False
+        Just mv -> putMVar mv payload >> return True
+
+  NMCast uuid payload -> do
+    rt <- ask
+    rid <- remoteNodeId
+    liftIO $ do
+      actors <- readTVarIO (rtActors rt)
+      case findByUUID uuid actors of
+        Nothing -> do
+          putStrLn $ "didn't find actor: " <> show uuid
+          return False
+        Just (SomeActorRef (LocalRef {arMsgQ})) ->
+          case decodeOrFail payload of
+            Left _            -> return False
+            Right (_, _, msg) -> do
+              atomically $ writeTQueue arMsgQ (rid, Cast msg)
+              return True
+        Just (SomeActorRef (RemoteRef _)) ->
+          return False
+
+  NMCall uuid corrId returnAddr payload -> do
+    rt <- ask
+    rid <- remoteNodeId
+    liftIO $ do
+      actors <- readTVarIO (rtActors rt)
+      case findByUUID uuid actors of
+        Nothing ->
+          return False
+        Just (SomeActorRef (LocalRef {arMsgQ})) ->
+          case decodeOrFail payload of
+            Left _            -> return False
+            Right (_, _, msg) -> do
+              mv <- newEmptyMVar
+              atomically $ writeTQueue arMsgQ (rid, Call msg mv)
+              void $ forkIO $ do
+                -- Bounded wait: if the target actor dies before replying,
+                -- this thread must not linger forever.
+                reply <- timeout 30_000_000 (takeMVar mv)
+                case reply of
+                  Nothing        -> return ()
+                  Just Nothing   -> return ()
+                  Just (Just rv) -> void $ try @SomeException $
+                    withRuntime rt $ do
+                      connRef <- getOrCreateConn returnAddr
+                      cast' (NMReply corrId (encode rv)) connRef
+              return True
+        Just (SomeActorRef (RemoteRef _)) ->
+          return False
+
+  NMDeath deadId reason -> do
+    routeRemoteDeath senderAddr deadId reason
+    return True
+
+  where
+    remoteNodeId = fromMaybe thisNodeId <$> addrToNodeId senderAddr
+
+-- System event handler
+
+sysHandlerFn :: Handler NetworkMessage () ()
+sysHandlerFn _ = return $ ActorResult Nothing () Nothing
+
+-- Node connection
+
+-- | Connect to a remote node and return its locally-assigned NodeId.
+-- NodeId 0 is always self; remote nodes get ids starting from 1.
+-- A suggested id is honoured if it is free (non-zero, not already in use).
+connect :: Maybe NodeId -> NodeAddr -> RuntimeM NodeId
+connect suggestedId peer = do
+  rt     <- ask
+  nodeId <- liftIO $ atomically $ do
+    table <- readTVar (rtNodeTable rt)
+    case Map.foldrWithKey (\k v acc -> if v == peer then Just k else acc) Nothing table of
+      Just existing -> return existing
+      Nothing -> do
+        nid <- readTVar (rtNextNodeId rt)
+        let assigned = case suggestedId of
+              Just n | n /= 0 && not (Map.member n table) -> n
+              _else                                        -> nid
+        writeTVar  (rtNextNodeId rt) (max (assigned + 1) nid)
+        modifyTVar (rtNodeTable rt) (Map.insert assigned peer)
+        return assigned
+  void $ getOrCreateConn peer
+  return nodeId

src/Control/Actor/Registry.hs

@@ -0,0 +1,141 @@
+module Control.Actor.Registry
+  ( registerActor,
+    lookupActor,
+    lookupRemoteActor,
+    registryUUID,
+    registry,
+    registerDeath,
+    deregisterNode,
+    createRegistry,
+  )
+where
+
+import Control.Actor.Core (ActorM, ActorResult (..), Handler, call, call', cast', liftRuntime, linkActorTo, pass, spawnActorAs, state, lastMessageFrom, cast, passWith)
+import Control.Actor.Runtime (Runtime (..), RuntimeM, getActorByUUID, withRuntime)
+import Control.Actor.Types (ActorId (..), ActorRef (..), DeathMessage (..), DeathTarget (LocalTarget), NodeId, RegistryMsg (..), SomeActorRef (..), SupervisorAction (Continue), actorRefId, thisNodeId, findByUUID)
+import Control.Concurrent.STM (readTVarIO)
+import Control.Monad (forM_, when)
+import Control.Monad.IO.Class (MonadIO (liftIO))
+import Control.Monad.Reader (ask)
+import Data.Map qualified as Map
+import Data.UUID (UUID, fromWords)
+import Unsafe.Coerce (unsafeCoerce)
+
+registryUUID :: UUID
+registryUUID = fromWords 0 0 0 1
+
+type RegistryState = Map.Map String ActorId
+
+registryHandlerFn :: Handler RegistryMsg RegistryState (Maybe ActorId)
+registryHandlerFn msg@(RMRegister name uuid) = do
+  u <- state
+  from <- lastMessageFrom
+  when (from == thisNodeId) $ do
+    rt <- liftRuntime ask
+    nodeTable <- liftIO $ readTVarIO (rtNodeTable rt)
+    forM_ (Map.keys nodeTable) $ \nodeId -> do
+      cast msg (RemoteRef (ActorId nodeId registryUUID))
+  passWith (Map.insert name (ActorId from uuid) u)
+registryHandlerFn (RMLookup name) = do
+  u <- state
+  return $ ActorResult (Just (Map.lookup name u)) u Nothing
+registryHandlerFn (RMDeregister name) = do
+  u <- state
+  return $ ActorResult Nothing (Map.delete name u) Nothing
+registryHandlerFn (RMDeregisterNode nodeId) = do
+  u <- state
+  passWith (Map.filter (\(ActorId nid _) -> nid /= nodeId) u)
+registryHandlerFn msg@(RMDeath deadUUID) = do
+  u <- state
+  from <- lastMessageFrom
+
+  let (_, kept) = Map.partition (\(ActorId _ uuid) -> uuid == deadUUID) u
+
+  when (from == thisNodeId) $ do
+    rt <- liftRuntime ask
+    nodeTable <- liftIO $ readTVarIO (rtNodeTable rt)
+    forM_ (Map.keys nodeTable) $ \nodeId -> do
+      cast msg (RemoteRef (ActorId nodeId registryUUID))
+
+  passWith kept
+
+registryDeathFn :: DeathMessage -> ActorM RegistryState (SupervisorAction RegistryState)
+registryDeathFn (DeathMessage (ActorId _ deadUUID) _) = do
+  u <- state
+  let (_, kept) = Map.partition (\(ActorId _ uuid) -> uuid == deadUUID) u
+  rt <- liftRuntime ask
+  nodeTable <- liftIO $ readTVarIO (rtNodeTable rt)
+  forM_ (Map.keys nodeTable) $ \nodeId -> do
+    cast (RMDeath deadUUID) (RemoteRef (ActorId nodeId registryUUID))
+
+  return $ Continue kept
+
+registry :: RuntimeM (Maybe (ActorRef RegistryMsg (Maybe ActorId)))
+registry = do
+  mRef <- getActorByUUID registryUUID
+  return $ case mRef of
+    Nothing                    -> Nothing
+    Just (SomeActorRef regRef) -> Just (unsafeCoerce regRef)
+
+registerDeath :: ActorId -> RuntimeM ()
+registerDeath (ActorId _ uuid) = do
+  r <- registry
+  mapM_ (cast' (RMDeath uuid)) r
+
+deregisterNode :: NodeId -> RuntimeM ()
+deregisterNode nodeId = do
+  r <- registry
+  mapM_ (cast' (RMDeregisterNode nodeId)) r
+
+registerActor :: String -> ActorRef msg r -> RuntimeM ()
+registerActor name ref = do
+  rt <- ask
+  let ActorId _ uuid = actorRefId ref
+  actors <- liftIO $ readTVarIO (rtActors rt)
+  case findByUUID registryUUID actors of
+    Just (SomeActorRef regRef) -> do
+      let reg = unsafeCoerce regRef :: ActorRef RegistryMsg (Maybe ActorId)
+      cast' (RMRegister name uuid) reg
+      case reg of
+        LocalRef {arDeathQ} -> linkActorTo (LocalTarget arDeathQ) ref
+        RemoteRef _ -> return ()
+    Nothing -> liftIO $ putStrLn "registerActor: registry not found"
+
+lookupActor :: String -> RuntimeM (Maybe (ActorRef msg r))
+lookupActor name = do
+  rt <- ask
+  actors <- liftIO $ readTVarIO (rtActors rt)
+  case findByUUID registryUUID actors of
+    Nothing -> return Nothing
+    Just (SomeActorRef regRef) -> do
+      result <-
+        call'
+          (RMLookup name)
+          (unsafeCoerce regRef :: ActorRef RegistryMsg (Maybe ActorId))
+      case result of
+        Just (Just (ActorId 0 uuid)) -> do
+          actors' <- liftIO $ readTVarIO (rtActors rt)
+          return $ fmap (\(SomeActorRef r) -> unsafeCoerce r) (findByUUID uuid actors')
+        Just (Just aid) -> return $ Just (RemoteRef aid)
+        _else -> return Nothing
+
+lookupRemoteActor :: String -> RuntimeM (Maybe (ActorRef msg r))
+lookupRemoteActor name = do
+  rt <- ask
+  table <- liftIO $ readTVarIO (rtNodeTable rt)
+  go (Map.keys table)
+  where
+    go [] = return Nothing
+    go (peerNodeId : rest) = do
+      let remoteReg =
+            RemoteRef (ActorId peerNodeId registryUUID) ::
+              ActorRef RegistryMsg (Maybe ActorId)
+      result <- call' (RMLookup name) remoteReg
+      case result of
+        Just (Just (ActorId 0 uuid)) ->
+          return $ Just (RemoteRef (ActorId peerNodeId uuid))
+        _else -> go rest
+
+createRegistry :: RuntimeM (ActorRef RegistryMsg (Maybe ActorId))
+createRegistry =
+  spawnActorAs registryUUID registryHandlerFn registryDeathFn Map.empty

src/Control/Actor/Runtime.hs

@@ -0,0 +1,103 @@
+{-# LANGUAGE StrictData #-}
+
+module Control.Actor.Runtime
+  ( Runtime (..)
+  , RuntimeM (..)
+  , newRuntime
+  , withRuntime
+  , lookupNode
+  , addrToNodeId
+  , getActorRef
+  , getActorByUUID
+  ) where
+
+import Control.Actor.Transport (Transport)
+import Control.Actor.Types
+import Control.Concurrent (ThreadId)
+import Control.Exception (SomeException)
+import Control.Concurrent.MVar (MVar)
+import Control.Concurrent.STM
+  ( TMVar, TVar, atomically, newTVarIO, readTVar, readTVarIO )
+import Control.Monad.Reader
+  ( MonadIO (..), MonadReader (..), ReaderT (..), runReaderT, asks )
+import Data.ByteString.Lazy (ByteString)
+import Data.Foldable (foldl')
+import Data.Map qualified as Map
+import Data.UUID (UUID)
+
+data Runtime = Runtime
+  { rtNodeId       :: NodeAddr
+  , rtNextNodeId   :: TVar NodeId
+  , rtActors       :: TVar (Map.Map ActorId (ThreadId, SomeActorRef))
+  , rtPending      :: TVar (Map.Map CorrelationId (MVar ByteString))
+  , rtNextCorr     :: TVar CorrelationId
+  , rtNodeTable    :: TVar (Map.Map NodeId NodeAddr)
+  , rtTransport    :: Transport
+  , rtConnections  :: TVar (Map.Map NodeAddr (ActorRef NetworkMessage ()))
+  , rtConnPromises :: TVar (Map.Map NodeAddr (TMVar (Either SomeException (ActorRef NetworkMessage ()))))
+  , rtSendRemote   :: NodeAddr -> NetworkMessage -> RuntimeM ()
+  }
+
+newtype RuntimeM a = RuntimeM
+  { unRuntimeM :: ReaderT Runtime IO a }
+  deriving (Functor, Applicative, Monad, MonadIO, MonadReader Runtime)
+
+newRuntime :: NodeAddr -> Transport -> IO Runtime
+newRuntime myAddr transport = do
+  actors    <- newTVarIO Map.empty
+  pending   <- newTVarIO Map.empty
+  nextCorr  <- newTVarIO (0 :: Integer)
+  nodeTable <- newTVarIO Map.empty
+  nextNid   <- newTVarIO (1 :: NodeId)
+  conns     <- newTVarIO Map.empty
+  promises  <- newTVarIO Map.empty
+  return Runtime
+    { rtNodeId       = myAddr
+    , rtNextNodeId   = nextNid
+    , rtActors       = actors
+    , rtPending      = pending
+    , rtNextCorr     = nextCorr
+    , rtNodeTable    = nodeTable
+    , rtTransport    = transport
+    , rtConnections  = conns
+    , rtConnPromises = promises
+    , rtSendRemote   = \_ _ -> return ()
+    }
+
+withRuntime :: Runtime -> RuntimeM a -> IO a
+withRuntime rt m = runReaderT (unRuntimeM m) rt
+{-# INLINE withRuntime #-}
+
+lookupNode :: NodeId -> RuntimeM (Maybe NodeAddr)
+lookupNode nodeId = do
+  rt <- ask
+  liftIO $ atomically $ do
+    table <- readTVar (rtNodeTable rt)
+    return $ Map.lookup nodeId table
+
+addrToNodeId :: NodeAddr -> RuntimeM (Maybe NodeId)
+addrToNodeId addr = do
+  rt <- ask
+  liftIO $ atomically $ do
+    table <- readTVar (rtNodeTable rt)
+    return $ Map.foldrWithKey (\k v acc -> if v == addr then Just k else acc) Nothing table
+
+getActorRef :: ActorId -> RuntimeM (Maybe SomeActorRef)
+getActorRef aid = do
+  actVar <- asks rtActors
+  actors <- liftIO $ readTVarIO actVar
+  return $ snd <$> Map.lookup aid actors
+
+firstByKey' :: Ord k => (k -> Bool) -> Map.Map k v -> Maybe v
+firstByKey' p = snd . foldl' step (False, Nothing) . Map.toAscList
+  where
+    step (True, mv) _ = (True, mv)
+    step (False, _) (k,v)
+      | p k       = (True, Just v)
+      | otherwise = (False, Nothing)
+
+getActorByUUID :: UUID -> RuntimeM (Maybe SomeActorRef)
+getActorByUUID uuid = do
+  actVar <- asks rtActors
+  actors <- liftIO $ readTVarIO actVar
+  return $ snd <$> firstByKey' (\(ActorId _ x) -> x == uuid) actors

src/Control/Actor/Supervision.hs

@@ -0,0 +1,190 @@
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE StrictData #-}
+
+module Control.Actor.Supervision
+  ( ChildSpec (..)
+  , child
+  , childWithRef
+  , RestartStrategy (..)
+  , ChildSlot (..)
+  , spawnSlot
+  , supervise'
+  , supervise
+  , doOneForOne
+  , doOneForAll
+  , doRestForOne
+  , killSlot
+  ) where
+
+import Control.Actor.Core (ActorM, Handler, liftRuntime, linkActorTo, spawnActor)
+import Control.Actor.Runtime (Runtime (..), RuntimeM, withRuntime)
+import Control.Actor.Types
+import Data.Binary (Binary)
+import Control.Concurrent (ThreadId, forkIO, killThread)
+import Control.Concurrent.STM
+  ( TMVar
+  , TVar
+  , atomically
+  , modifyTVar
+  , newTQueueIO
+  , newTVarIO
+  , putTMVar
+  , readTQueue
+  , readTVarIO
+  , tryTakeTMVar
+  , writeTVar
+  )
+import Control.Monad (void)
+import GHC.Clock (getMonotonicTime)
+import Control.Monad.Reader (MonadIO (..), MonadReader (..))
+import Data.Map qualified as Map
+
+data ChildSpec = forall m r. (Binary m, Binary r) => ChildSpec
+  { csRun     :: DeathTarget -> RuntimeM (ActorRef m r)
+  -- | Runs after every (re)spawn — including supervisor restarts — so it can
+  -- refresh shared ref cells and re-register names.
+  , csOnSpawn :: ActorRef m r -> RuntimeM ()
+  }
+
+child ::
+  (Binary m, Binary r) =>
+  Handler m u r ->
+  (DeathMessage -> ActorM u (SupervisorAction u)) ->
+  u ->
+  ChildSpec
+child msgFn deathFn initState =
+  ChildSpec
+    { csRun = \target -> do
+        ref <- spawnActor msgFn deathFn initState
+        linkActorTo target ref
+        return ref
+    , csOnSpawn = \_ -> return ()
+    }
+
+childWithRef ::
+  (Binary m, Binary r) =>
+  Handler m u r ->
+  (DeathMessage -> ActorM u (SupervisorAction u)) ->
+  u ->
+  TMVar (ActorRef m r) ->
+  ChildSpec
+childWithRef msgFn deathFn initState cell =
+  ChildSpec
+    { csRun = \target -> do
+        ref <- spawnActor msgFn deathFn initState
+        linkActorTo target ref
+        return ref
+    , csOnSpawn = \ref -> liftIO $ atomically $ do
+        void $ tryTakeTMVar cell
+        putTMVar cell ref
+    }
+
+data RestartStrategy = OneForOne | OneForAll | RestForOne
+
+data ChildSlot = forall m r. (Binary m, Binary r) => ChildSlot
+  { slotSpec :: ChildSpec
+  , slotRef  :: ActorRef m r
+  , slotId   :: ActorId
+  , slotTid  :: ThreadId
+  }
+
+spawnSlot :: DeathTarget -> ChildSpec -> RuntimeM ChildSlot
+spawnSlot target spec@ChildSpec {csRun, csOnSpawn} = do
+  rt  <- ask
+  ref <- csRun target
+  csOnSpawn ref
+  let aid = someActorId (SomeActorRef ref)
+  tid <- liftIO $ do
+    actors <- readTVarIO (rtActors rt)
+    case Map.lookup aid actors of
+      Just (t, _) -> return t
+      Nothing     -> error "spawnSlot: actor vanished immediately after spawn"
+  return $ ChildSlot spec ref aid tid
+
+killSlot :: ChildSlot -> IO ()
+killSlot ChildSlot {slotTid} = killThread slotTid
+
+-- | Max child restarts within 'restartWindow' seconds before the supervisor
+-- gives up and stops all children — otherwise a child that crashes right
+-- after spawning would respawn in a hot loop forever.
+restartLimit :: Int
+restartLimit = 5
+
+restartWindow :: Double
+restartWindow = 10
+
+supervise' :: RestartStrategy -> [ChildSpec] -> RuntimeM ()
+supervise' strategy specs = do
+  rt        <- ask
+  supDeathQ <- liftIO newTQueueIO
+  let target = LocalTarget supDeathQ
+  slots    <- mapM (spawnSlot target) specs
+  slotsVar <- liftIO $ newTVarIO slots
+  let loop restarts = do
+        DeathMessage deadId reason <- atomically $ readTQueue supDeathQ
+        slots' <- readTVarIO slotsVar
+        if all ((/= deadId) . slotId) slots'
+          -- Stale death: killSlot is asynchronous, so children replaced by a
+          -- OneForAll/RestForOne restart report their deaths after the fact.
+          -- Acting on one would kill the fresh children and cascade forever.
+          then loop restarts
+          else case reason of
+            -- A normal exit is completion, not a crash: drop the slot.
+            Normal -> do
+              atomically $ modifyTVar slotsVar (filter ((/= deadId) . slotId))
+              loop restarts
+            _crash -> do
+              now <- getMonotonicTime
+              let recent = filter (> now - restartWindow) restarts
+              if length recent >= restartLimit
+                then do
+                  putStrLn $ "supervisor: " <> show restartLimit
+                    <> " restarts within " <> show restartWindow
+                    <> "s, giving up"
+                  mapM_ killSlot slots'
+                else do
+                  withRuntime rt $ case strategy of
+                    OneForOne  -> doOneForOne target slotsVar slots' deadId
+                    OneForAll  -> doOneForAll target slotsVar slots'
+                    RestForOne -> doRestForOne target slotsVar slots' deadId
+                  loop (now : recent)
+  liftIO $ void $ forkIO $ loop []
+
+supervise :: RestartStrategy -> [ChildSpec] -> ActorM u ()
+supervise = (liftRuntime .) . supervise'
+
+doOneForOne ::
+  DeathTarget -> TVar [ChildSlot] -> [ChildSlot] -> ActorId -> RuntimeM ()
+doOneForOne target slotsVar slots deadId =
+  case filter (\s -> slotId s == deadId) slots of
+    []     -> return ()
+    slot:_ -> do
+      newSlot <- spawnSlot target (slotSpec slot)
+      liftIO $ atomically $
+        modifyTVar slotsVar $
+          map (\s -> if slotId s == deadId then newSlot else s)
+
+doOneForAll ::
+  DeathTarget ->
+  TVar [ChildSlot] ->
+  [ChildSlot] ->
+  RuntimeM ()
+doOneForAll target slotsVar slots = do
+  liftIO $ mapM_ killSlot slots
+  newSlots <- mapM (spawnSlot target . slotSpec) slots
+  liftIO $ atomically $ writeTVar slotsVar newSlots
+
+doRestForOne ::
+  DeathTarget ->
+  TVar [ChildSlot] ->
+  [ChildSlot] ->
+  ActorId ->
+  RuntimeM ()
+doRestForOne target slotsVar slots deadId = do
+  let (before, fromDead) = break (\s -> slotId s == deadId) slots
+  case fromDead of
+    []        -> return ()
+    _nonempty -> do
+      liftIO $ mapM_ killSlot (drop 1 fromDead)
+      newSlots <- mapM (spawnSlot target . slotSpec) fromDead
+      liftIO $ atomically $ writeTVar slotsVar (before ++ newSlots)

src/Control/Actor/Transport.hs

@@ -0,0 +1,112 @@
+{-# LANGUAGE StrictData #-}
+
+module Control.Actor.Transport
+  ( ConnHandle (..)
+  , Transport (..)
+  , createTCPTransport
+  ) where
+
+import Control.Actor.Types (NodeAddr (..))
+import Control.Concurrent (forkIO)
+import Control.Monad (forever, void, when)
+import Data.Binary (decode, encode)
+import qualified Data.ByteString.Lazy as BL
+import Data.ByteString.Lazy (ByteString)
+import Data.Word (Word64)
+import Network.Socket
+  ( AddrInfo (addrAddress, addrFamily, addrFlags, addrProtocol, addrSocketType)
+  , AddrInfoFlag (AI_PASSIVE)
+  , Socket
+  , SocketOption (NoDelay, ReuseAddr)
+  , SocketType (Stream)
+  , accept
+  , bind
+  , close
+  , connect
+  , defaultHints
+  , getAddrInfo
+  , listen
+  , setSocketOption
+  , socket
+  , socketPort
+  )
+import Network.Socket.ByteString.Lazy (recv, sendAll)
+
+data ConnHandle = ConnHandle
+  { chSend  :: ByteString -> IO ()
+  , chRecv  :: IO ByteString
+  , chClose :: IO ()
+  }
+
+data Transport = Transport
+  { tConnect :: NodeAddr -> IO ConnHandle
+  , tListen  :: (ConnHandle -> IO ()) -> IO ()
+  }
+
+sendFramed :: Socket -> ByteString -> IO ()
+sendFramed sock payload =
+  sendAll sock (encode (fromIntegral (BL.length payload) :: Word64) <> payload)
+
+recvExact :: Socket -> Int -> IO ByteString
+recvExact sock = go []
+  where
+    go acc 0 = return (BL.concat (reverse acc))
+    go acc r = do
+      chunk <- recv sock (fromIntegral (min r 65536))
+      when (BL.null chunk) $ fail "recvExact: connection closed"
+      go (chunk : acc) (r - fromIntegral (BL.length chunk))
+
+recvFramed :: Socket -> IO ByteString
+recvFramed sock = do
+  header <- recvExact sock 8
+  recvExact sock (fromIntegral (decode header :: Word64))
+
+connectTcp :: NodeAddr -> IO Socket
+connectTcp (NodeAddr host port) = do
+  let hints = defaultHints {addrSocketType = Stream}
+  addrs <- getAddrInfo (Just hints) (Just host) (Just (show port))
+  case addrs of
+    []    -> fail $ "connectTcp: no address for " <> host <> ":" <> show port
+    (a:_) -> do
+      sock <- socket (addrFamily a) Stream (addrProtocol a)
+      setSocketOption sock NoDelay 1
+      connect sock (addrAddress a)
+      return sock
+
+listenTcp :: NodeAddr -> IO (Socket, NodeAddr)
+listenTcp (NodeAddr host port) = do
+  let hints = defaultHints {addrFlags = [AI_PASSIVE], addrSocketType = Stream}
+  addrs <- getAddrInfo (Just hints) (Just host) (Just (show port))
+  case addrs of
+    []    -> fail "listenTcp: no address"
+    (a:_) -> do
+      sock <- socket (addrFamily a) Stream (addrProtocol a)
+      setSocketOption sock ReuseAddr 1
+      bind sock (addrAddress a)
+      listen sock 128
+      actualPort <- fromIntegral <$> socketPort sock
+      return (sock, NodeAddr host actualPort)
+
+-- | Create a TCP transport bound to the given address.
+-- Pass port 0 to let the OS pick a free port.
+-- Returns the transport and the address actually bound (with the real port).
+createTCPTransport :: NodeAddr -> IO (Transport, NodeAddr)
+createTCPTransport myAddr = do
+  (lsock, actualAddr) <- listenTcp myAddr
+  let transport = Transport
+        { tConnect = \peer -> do
+            sock <- connectTcp peer
+            return ConnHandle
+              { chSend  = sendFramed sock
+              , chRecv  = recvFramed sock
+              , chClose = close sock
+              }
+        , tListen = \callback -> void $ forkIO $ forever $ do
+            (csock, _) <- accept lsock
+            void $ forkIO $ callback ConnHandle
+              { chSend  = sendFramed csock
+              , chRecv  = recvFramed csock
+              , chClose = close csock
+              }
+        }
+  return (transport, actualAddr)

src/Control/Actor/Types.hs

@@ -0,0 +1,141 @@
+{-# LANGUAGE ExistentialQuantification #-}
+{-# LANGUAGE StrictData #-}
+
+module Control.Actor.Types
+  ( NodeAddr (..)
+  , NodeId
+  , thisNodeId
+  , ActorId (..)
+  , ExitReason (..)
+  , DeathMessage (..)
+  , DeathTarget (..)
+  , RemoteExitReason (..)
+  , toRemoteExitReason
+  , NetworkMessage (..)
+  , ActorState (..)
+  , ActorRef (..)
+  , SomeActorRef (..)
+  , someActorId
+  , actorRefId
+  , SupervisorAction (..)
+  , CorrelationId
+  , Envelope (..)
+  , RegistryMsg (..)
+  , findByUUID
+  ) where
+
+import Control.Concurrent.MVar (MVar)
+import Control.Concurrent.STM (TQueue, TVar)
+import Control.Exception (SomeException)
+import Data.Binary (Binary)
+import Data.ByteString.Lazy (ByteString)
+import Data.UUID (UUID)
+import GHC.Generics (Generic)
+import Data.Map qualified as Map
+import Control.Concurrent (ThreadId)
+
+data NodeAddr = NodeAddr
+  { nodeHost :: String
+  , nodePort :: Integer
+  } deriving (Eq, Ord, Show, Generic)
+
+instance Binary NodeAddr
+
+type NodeId = Integer
+
+thisNodeId :: NodeId
+thisNodeId = 0
+
+data ActorId = ActorId NodeId UUID
+  deriving (Eq, Ord, Show, Generic)
+
+instance Binary ActorId
+
+data ExitReason
+  = Normal
+  | Killed
+  | Exception SomeException
+  deriving Show
+
+data DeathMessage = DeathMessage
+  { dmActorId :: ActorId
+  , dmReason  :: ExitReason
+  } deriving Show
+
+data DeathTarget
+  = LocalTarget (TQueue DeathMessage)
+  | RemoteTarget ActorId NodeId
+
+data RemoteExitReason
+  = RNormal
+  | RKilled
+  | RException String
+  deriving (Show, Generic)
+
+instance Binary RemoteExitReason
+
+toRemoteExitReason :: ExitReason -> RemoteExitReason
+toRemoteExitReason Normal        = RNormal
+toRemoteExitReason Killed        = RKilled
+toRemoteExitReason (Exception e) = RException (show e)
+
+data NetworkMessage
+  = NMHandshake NodeAddr
+  | NMCast UUID ByteString
+  | NMCall UUID CorrelationId NodeAddr ByteString
+  | NMReply CorrelationId ByteString
+  | NMDeath ActorId RemoteExitReason
+  deriving (Generic, Show)
+
+instance Binary NetworkMessage
+
+data ActorState u = ActorState
+  { asId    :: ActorId
+  , asLinks :: TVar [DeathTarget]
+  , asEnv   :: u
+  , asLatestFrom :: NodeId
+  }
+
+type CorrelationId = Integer
+
+data Envelope msg reply
+  = Cast msg
+  | Call msg (MVar (Maybe reply))
+
+data ActorRef msg reply
+  = forall u. LocalRef
+      { arMsgQ   :: TQueue (NodeId, Envelope msg reply)
+      , arDeathQ :: TQueue DeathMessage
+      , arState  :: ActorState u
+      }
+  | RemoteRef ActorId
+
+data SomeActorRef = forall msg reply. (Binary msg, Binary reply) => SomeActorRef (ActorRef msg reply)
+
+someActorId :: SomeActorRef -> ActorId
+someActorId (SomeActorRef (LocalRef {arState})) = asId arState
+someActorId (SomeActorRef (RemoteRef aid))      = aid
+
+actorRefId :: ActorRef msg r -> ActorId
+actorRefId (LocalRef {arState}) = asId arState
+actorRefId (RemoteRef aid)      = aid
+
+data SupervisorAction u
+  = Stop
+  | Continue u
+
+data RegistryMsg
+  = RMRegister        String UUID
+  | RMLookup          String
+  | RMDeregister      String
+  | RMDeath           UUID
+  | RMDeregisterNode  NodeId
+  deriving (Generic)
+
+instance Binary RegistryMsg
+
+findByUUID :: UUID -> Map.Map ActorId (ThreadId, SomeActorRef) -> Maybe SomeActorRef
+findByUUID uuid actors =
+  case Map.toList (Map.filterWithKey (\(ActorId _ u) _ -> u == uuid) actors) of
+    []            -> Nothing
+    (_, (_, r)):_ -> Just r

src/FL/Actors/Aggregator.hs

@@ -0,0 +1,43 @@
+module FL.Actors.Aggregator
+  ( AggregatorState (..)
+  , aggregatorHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Control.Concurrent.STM (TMVar, atomically, readTMVar)
+import Control.Monad.IO.Class (liftIO)
+import Data.List (sortOn)
+import FL.Actors.Logger (logFn)
+import FL.NN (fedAvg)
+import FL.Types
+
+data AggregatorState = AggregatorState
+  { asExpected     :: Int
+  , asUpdates      :: [ModelUpdate]
+  , asCoordCell    :: TMVar (ActorRef CoordinatorMsg ())  -- filled after coordinator spawns
+  , asLogRef       :: ActorRef LogMsg ()
+  , asCurrentRound :: Int
+  }
+
+aggregatorHandler :: Handler AggregatorMsg AggregatorState ()
+aggregatorHandler (SubmitUpdate upd) = do
+  st <- state
+  if muRound upd /= asCurrentRound st
+    then pass
+    else do
+      let updates' = upd : filter (\u -> muTrainerId u /= muTrainerId upd) (asUpdates st)
+      if length updates' >= asExpected st
+        then do
+          let pairs    = [(muSampleSize u, wData (muWeights u)) | u <- sortOn muTrainerId updates']
+              avgW     = Weights (fedAvg pairs)
+              nextRound = asCurrentRound st + 1
+              newModel  = GlobalModel nextRound avgW
+          logFn (asLogRef st) "Aggregator" INFO $
+            "Round " <> show (asCurrentRound st)
+            <> ": aggregated " <> show (length updates') <> " updates"
+          coordRef <- liftIO $ atomically $ readTMVar (asCoordCell st)
+          cast (AggregationComplete newModel) coordRef
+          passWith st { asUpdates = [], asCurrentRound = nextRound }
+        else
+          passWith st { asUpdates = updates' }

src/FL/Actors/Coordinator.hs

@@ -0,0 +1,103 @@
+module FL.Actors.Coordinator
+  ( CoordinatorState (..)
+  , coordinatorHandler
+  , coordinatorDeathHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Data.Binary (encodeFile)
+import Control.Concurrent.MVar (MVar, putMVar)
+import Control.Monad (forM_, when)
+import Control.Monad.IO.Class (liftIO)
+import Data.Map.Strict (Map)
+import Data.Map.Strict qualified as Map
+import Data.Set qualified as Set
+import Data.UUID.V4 (nextRandom)
+import FL.Actors.Logger (logFn)
+import FL.CRDT (ORSet, orAdd, orMembers)
+import FL.Types
+
+data CoordinatorState = CoordinatorState
+  { csTrainers         :: ORSet String
+  , csTrainerRefs      :: Map String (ActorRef TrainerMsg ())
+  , csExpectedTrainers :: Int
+  , csCurrentRound     :: Int
+  , csTotalRounds      :: Int
+  , csGlobalModel      :: GlobalModel
+  , csAggRef           :: ActorRef AggregatorMsg ()
+  , csEvalRef          :: ActorRef EvaluatorMsg ()
+  , csLogRef           :: ActorRef LogMsg ()
+  , csDoneVar          :: MVar ()
+  }
+
+coordinatorHandler :: Handler CoordinatorMsg CoordinatorState ()
+coordinatorHandler (RegisterTrainer tid aid) = do
+  st <- state
+  -- The ActorId in the message carries the *sender's* node numbering (its own
+  -- node id is always 0 to itself); rebase it onto our node id for the sender,
+  -- otherwise every cast to the trainer is dropped ("no node with id 0").
+  from <- lastMessageFrom
+  uuid <- liftIO nextRandom
+  let ActorId _ trainerUUID = aid
+      ref       = RemoteRef (ActorId from trainerUUID) :: ActorRef TrainerMsg ()
+      trainers' = orAdd uuid tid (csTrainers st)
+      refs'     = Map.insert tid ref (csTrainerRefs st)
+      st'       = st { csTrainers = trainers', csTrainerRefs = refs' }
+      numRegistered = Set.size (orMembers trainers')
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Trainer registered: " <> tid
+    <> " (" <> show numRegistered <> "/" <> show (csExpectedTrainers st) <> ")"
+  -- Auto-start once all expected trainers have registered (per the OR-Set)
+  if numRegistered >= csExpectedTrainers st
+    then startRound st'
+    else passWith st'
+
+coordinatorHandler StartFederation = do
+  st <- state
+  startRound st
+
+coordinatorHandler (AggregationComplete newModel) = do
+  st <- state
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Round " <> show (gmRound newModel - 1) <> " aggregation complete"
+  cast (EvaluateModel newModel) (csEvalRef st)
+  let st' = st { csGlobalModel = newModel, csCurrentRound = gmRound newModel }
+  if csCurrentRound st' >= csTotalRounds st'
+    then do
+      liftIO $ encodeFile "fl_model_coordinator.bin" (gmWeights newModel)
+      logFn (csLogRef st') "Coordinator" INFO
+        "Model saved to fl_model_coordinator.bin"
+      passWith st'
+    else
+      startRound st'
+
+coordinatorHandler (EvaluationDone metrics) = do
+  st <- state
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Eval round=" <> show (mRound metrics)
+    <> " acc="  <> show4 (mAccuracy  metrics)
+    <> " prec=" <> show4 (mPrecision metrics)
+    <> " rec="  <> show4 (mRecall    metrics)
+    <> " f1="   <> show4 (mF1        metrics)
+    <> " auc="  <> show4 (mAUCROC    metrics)
+  liftIO $ when (mRound metrics >= csTotalRounds st) $
+    putMVar (csDoneVar st) ()
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
+
+startRound :: CoordinatorState -> Actor CoordinatorMsg CoordinatorState ()
+startRound st = do
+  let model = csGlobalModel st
+  logFn (csLogRef st) "Coordinator" INFO $
+    "Starting round " <> show (gmRound model + 1)
+    <> " with " <> show (Map.size (csTrainerRefs st)) <> " trainers"
+  forM_ (Map.elems (csTrainerRefs st)) $ \ref ->
+    cast (StartTraining model) ref
+  passWith st
+
+coordinatorDeathHandler :: DeathMessage -> ActorM CoordinatorState (SupervisorAction CoordinatorState)
+coordinatorDeathHandler _ = do
+  st <- state
+  return (Continue st)

src/FL/Actors/Evaluator.hs

@@ -0,0 +1,36 @@
+module FL.Actors.Evaluator
+  ( EvaluatorState (..)
+  , evaluatorHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Control.Concurrent.STM (TMVar, atomically, readTMVar)
+import Control.Monad.IO.Class (liftIO)
+import FL.Actors.Logger (logFn)
+import FL.NN (LocalData, evaluate, unflattenWeights)
+import FL.Types
+
+data EvaluatorState = EvaluatorState
+  { esTestData  :: LocalData
+  , esCoordCell :: TMVar (ActorRef CoordinatorMsg ())  -- filled after coordinator spawns
+  , esLogRef    :: ActorRef LogMsg ()
+  }
+
+evaluatorHandler :: Handler EvaluatorMsg EvaluatorState ()
+evaluatorHandler (EvaluateModel gm) = do
+  st <- state
+  let ws      = unflattenWeights (wData (gmWeights gm))
+      metrics = evaluate ws (esTestData st) (gmRound gm)
+  logFn (esLogRef st) "Evaluator" INFO $
+    "Round " <> show (gmRound gm)
+    <> "  acc="  <> show4 (mAccuracy  metrics)
+    <> "  prec=" <> show4 (mPrecision metrics)
+    <> "  rec="  <> show4 (mRecall    metrics)
+    <> "  f1="   <> show4 (mF1        metrics)
+    <> "  auc="  <> show4 (mAUCROC    metrics)
+  coordRef <- liftIO $ atomically $ readTMVar (esCoordCell st)
+  cast (EvaluationDone metrics) coordRef
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)

src/FL/Actors/Logger.hs

@@ -0,0 +1,37 @@
+module FL.Actors.Logger
+  ( loggerHandler
+  , logFn
+  , logFn'
+  ) where
+
+import Control.Actor.Core (ActorM, Handler, cast, cast', pass)
+import Control.Actor.Runtime (RuntimeM)
+import Control.Actor.Types
+import Control.Monad.IO.Class (liftIO)
+import Data.Time.Clock (getCurrentTime)
+import FL.Types
+
+type LogState = ()
+
+loggerHandler :: Handler LogMsg LogState ()
+loggerHandler (Log LogEntry{leSource, leLevel, leMsg}) = do
+  t <- liftIO getCurrentTime
+  liftIO $ putStrLn $
+    "[" <> show t <> "] [" <> showLevel leLevel <> "] "
+    <> leSource <> ": " <> leMsg
+  pass
+  where
+    showLevel DEBUG = "DEBUG"
+    showLevel INFO  = "INFO "
+    showLevel WARN  = "WARN "
+    showLevel ERROR = "ERROR"
+
+-- | Log via the logger actor. Goes through the normal cast path so it works
+-- for both local and remote logger refs (a remote ref used to be a silent
+-- no-op, losing all trainer-node logs in provider mode).
+logFn :: ActorRef LogMsg () -> String -> LogLevel -> String -> ActorM u ()
+logFn ref src lvl msg = cast (Log (LogEntry src lvl msg)) ref
+
+-- | Same as 'logFn' but usable outside an actor (node startup code).
+logFn' :: ActorRef LogMsg () -> String -> LogLevel -> String -> RuntimeM ()
+logFn' ref src lvl msg = cast' (Log (LogEntry src lvl msg)) ref

src/FL/Actors/PeerCoordinator.hs

@@ -0,0 +1,166 @@
+{-# LANGUAGE StrictData #-}
+
+module FL.Actors.PeerCoordinator
+  ( PeerCoordinatorState (..)
+  , peerCoordinatorHandler
+  , peerCoordinatorDeathHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Data.Binary (encodeFile)
+import Unsafe.Coerce (unsafeCoerce)
+import Control.Concurrent.MVar (MVar, putMVar)
+import Control.Monad (forM_, when)
+import Control.Monad.IO.Class (liftIO)
+import Data.List (sortOn)
+import Data.Set qualified as Set
+import Data.UUID.V4 (nextRandom)
+import FL.Actors.Logger (logFn)
+import FL.CRDT
+import FL.NN (fedAvg)
+import FL.Types
+
+data PeerCoordinatorState = PeerCoordinatorState
+  { pcsNodeId       :: String
+  , pcsRoundCounter :: GCounter        -- completed training rounds per peer (monotone, never reset)
+  , pcsParticipants :: ORSet String    -- active peers ever seen (monotone, never reset)
+  , pcsCurrentRound :: Int
+  , pcsTotalRounds  :: Int
+  , pcsGlobalModel  :: GlobalModel
+  , pcsPeerRefs     :: [ActorRef PeerCoordinatorMsg ()]
+  , pcsSelfRef      :: ActorRef PeerCoordinatorMsg ()
+  , pcsTrainerRef   :: ActorRef TrainerMsg ()
+  , pcsEvalRef      :: ActorRef EvaluatorMsg ()
+  , pcsLogRef       :: ActorRef LogMsg ()
+  , pcsCollected    :: [ModelUpdate]   -- updates for the current round + buffered future rounds
+  , pcsNumPeers     :: Int  -- total number of peers (including self)
+  , pcsDoneVar      :: MVar ()
+  }
+
+-- | Insert an update, deduplicating on (trainer, round). Stale (past-round)
+-- updates are dropped; future-round updates from fast peers are buffered so
+-- they are not lost — without the round key a fast peer's round-(r+1) update
+-- would overwrite its round-r update in a slower peer and rounds would mix.
+insertUpdate :: Int -> ModelUpdate -> [ModelUpdate] -> [ModelUpdate]
+insertUpdate currentRound upd collected
+  | muRound upd < currentRound = collected
+  | otherwise =
+      upd : filter (\u -> (muTrainerId u, muRound u) /= (muTrainerId upd, muRound upd)) collected
+
+currentRoundUpdates :: PeerCoordinatorState -> [ModelUpdate]
+currentRoundUpdates st = [u | u <- pcsCollected st, muRound u == pcsCurrentRound st]
+
+-- | A round is complete when the CRDTs agree with the collected updates:
+-- the OR-Set has seen all peers, the G-Counter shows every one of them has
+-- finished training the current round, and we hold all their updates.
+roundComplete :: PeerCoordinatorState -> Bool
+roundComplete st =
+  let r       = pcsCurrentRound st
+      members = orMembers (pcsParticipants st)
+      trained = Set.filter (\p -> gcGet p (pcsRoundCounter st) > r) members
+  in Set.size members >= pcsNumPeers st
+     && Set.size trained >= pcsNumPeers st
+     && length (currentRoundUpdates st) >= pcsNumPeers st
+
+peerCoordinatorHandler :: Handler PeerCoordinatorMsg PeerCoordinatorState ()
+peerCoordinatorHandler (SetPeers peerIds) = do
+  st <- state
+  (SomeActorRef self) <- getSelf
+  let peers   = map (\aid -> RemoteRef aid :: ActorRef PeerCoordinatorMsg ()) peerIds
+      selfRef = unsafeCoerce self :: ActorRef PeerCoordinatorMsg ()
+  passWith st { pcsPeerRefs = peers, pcsSelfRef = selfRef }
+
+peerCoordinatorHandler PeerStartRound = do
+  st <- state
+  logFn (pcsLogRef st) (pcsNodeId st) INFO $
+    "Starting local training for P2P round " <> show (pcsCurrentRound st + 1)
+  cast (StartTraining (pcsGlobalModel st)) (pcsTrainerRef st)
+  pass
+
+peerCoordinatorHandler (PeerTrainingDone upd) = do
+  st <- state
+  uuid <- liftIO nextRandom
+  let counter'      = gcIncrement (pcsNodeId st) (pcsRoundCounter st)
+      participants' = orAdd uuid (pcsNodeId st) (pcsParticipants st)
+      collected'    = insertUpdate (pcsCurrentRound st) upd (pcsCollected st)
+      payload       = PeerSyncPayload
+                        { pspCounter      = counter'
+                        , pspParticipants = participants'
+                        , pspUpdate       = Just upd
+                        }
+  logFn (pcsLogRef st) (pcsNodeId st) INFO
+    "Training done, broadcasting PeerSync"
+  forM_ (pcsPeerRefs st) $ \ref ->
+    cast (PeerSync payload) ref
+  let st' = st { pcsRoundCounter = counter'
+               , pcsParticipants = participants'
+               , pcsCollected    = collected'
+               }
+  if roundComplete st'
+    then finishRound st'
+    else passWith st'
+
+peerCoordinatorHandler (PeerSync payload) = do
+  st <- state
+  -- Merge CRDTs
+  let counter'      = gcMerge (pcsRoundCounter st) (pspCounter payload)
+      participants' = orMerge (pcsParticipants st) (pspParticipants payload)
+      collected'    = case pspUpdate payload of
+        Nothing  -> pcsCollected st
+        Just upd -> insertUpdate (pcsCurrentRound st) upd (pcsCollected st)
+      st' = st { pcsRoundCounter = counter'
+               , pcsParticipants = participants'
+               , pcsCollected    = collected'
+               }
+  if roundComplete st'
+    then finishRound st'
+    else passWith st'
+
+peerCoordinatorHandler (PeerEvaluationDone metrics) = do
+  st <- state
+  logFn (pcsLogRef st) (pcsNodeId st) INFO $
+    "Eval round=" <> show (mRound metrics)
+    <> " acc="  <> show4 (mAccuracy  metrics)
+    <> " auc="  <> show4 (mAUCROC    metrics)
+  liftIO $ when (mRound metrics >= pcsTotalRounds st) $
+    putMVar (pcsDoneVar st) ()
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
+
+finishRound :: PeerCoordinatorState -> Actor PeerCoordinatorMsg PeerCoordinatorState ()
+finishRound st = do
+  -- Sort so every peer folds the weighted average in the same order and
+  -- arrives at bit-identical models regardless of message arrival order.
+  let updates   = sortOn muTrainerId (currentRoundUpdates st)
+      pairs     = [(muSampleSize u, wData (muWeights u)) | u <- updates]
+      avgW      = Weights (fedAvg pairs)
+      nextRound = pcsCurrentRound st + 1
+      newModel  = GlobalModel nextRound avgW
+  logFn (pcsLogRef st) (pcsNodeId st) INFO $
+    "P2P round " <> show (pcsCurrentRound st) <> " aggregated ("
+    <> show (length updates) <> " updates)"
+  cast (EvaluateModel newModel) (pcsEvalRef st)
+  let st' = st
+        { pcsGlobalModel  = newModel
+        , pcsCurrentRound = nextRound
+        -- Keep buffered updates from peers already in later rounds
+        , pcsCollected    = [u | u <- pcsCollected st, muRound u >= nextRound]
+        }
+  if nextRound >= pcsTotalRounds st
+    then do
+      let filename = "fl_model_" <> pcsNodeId st' <> ".bin"
+      liftIO $ encodeFile filename (gmWeights newModel)
+      logFn (pcsLogRef st') (pcsNodeId st') INFO $
+        "Model saved to " <> filename
+      passWith st'
+    else do
+      -- Schedule next round after a short delay to let PeerSync propagate
+      castIn 500 PeerStartRound (pcsSelfRef st')
+      passWith st'
+
+peerCoordinatorDeathHandler :: DeathMessage -> ActorM PeerCoordinatorState (SupervisorAction PeerCoordinatorState)
+peerCoordinatorDeathHandler _ = do
+  st <- state
+  return (Continue st)

src/FL/Actors/Trainer.hs

@@ -0,0 +1,45 @@
+module FL.Actors.Trainer
+  ( TrainerState (..)
+  , trainerHandler
+  ) where
+
+import Control.Actor.Core
+import Control.Actor.Types
+import Control.Monad.IO.Class (liftIO)
+import FL.Actors.Logger (logFn)
+import FL.NN (LocalData, flattenWeights, localDataSize, trainLocal, unflattenWeights)
+import FL.Types
+
+data TrainerState = TrainerState
+  { tsId        :: String
+  , tsData      :: LocalData
+  , tsLastRound :: Int
+  , tsAggRef    :: ActorRef AggregatorMsg ()
+  , tsLogRef    :: ActorRef LogMsg ()
+  , tsConfig    :: FLConfig
+  }
+
+trainerHandler :: Handler TrainerMsg TrainerState ()
+trainerHandler (StartTraining gm) = do
+  st <- state
+  let round_ = gmRound gm
+  if round_ <= tsLastRound st
+    then pass  -- idempotent: already processed this round
+    else do
+      logFn (tsLogRef st) (tsId st) INFO $
+        "Starting local training for round " <> show round_
+      let initW = unflattenWeights (wData (gmWeights gm))
+      updatedW <- liftIO $ trainLocal (tsConfig st) (tsId st) round_ initW (tsData st)
+      let n      = localDataSize (tsData st)
+          update = ModelUpdate
+            { muTrainerId  = tsId st
+            , muRound      = round_
+            , muSampleSize = n
+            , muWeights    = Weights (flattenWeights updatedW)
+            }
+      logFn (tsLogRef st) (tsId st) INFO $
+        "Round " <> show round_ <> " training complete, samples=" <> show n
+      cast (SubmitUpdate update) (tsAggRef st)
+      passWith st { tsLastRound = round_ }
+
+trainerHandler TrainerStop = pass

src/FL/CRDT.hs

@@ -0,0 +1,82 @@
+module FL.CRDT
+  ( GCounter
+  , gcEmpty
+  , gcIncrement
+  , gcMerge
+  , gcValue
+  , gcGet
+  , ORSet
+  , orEmpty
+  , orAdd
+  , orRemove
+  , orMerge
+  , orMembers
+  ) where
+
+import Data.Binary (Binary (..))
+import Data.Map.Strict (Map)
+import Data.Map.Strict qualified as Map
+import Data.Set (Set)
+import Data.Set qualified as Set
+import Data.UUID (UUID)
+import GHC.Generics (Generic)
+
+-- ---------------------------------------------------------------------------
+-- G-Counter: each node has its own non-decreasing counter.
+-- Merge takes element-wise max. Value is the sum.
+-- ---------------------------------------------------------------------------
+
+newtype GCounter = GCounter (Map String Int)
+  deriving (Show, Eq, Generic)
+
+instance Binary GCounter where
+  put (GCounter m) = put (Map.toAscList m)
+  get = GCounter . Map.fromAscList <$> get
+
+gcEmpty :: GCounter
+gcEmpty = GCounter Map.empty
+
+gcIncrement :: String -> GCounter -> GCounter
+gcIncrement nodeId (GCounter m) =
+  GCounter (Map.insertWith (+) nodeId 1 m)
+
+gcMerge :: GCounter -> GCounter -> GCounter
+gcMerge (GCounter a) (GCounter b) =
+  GCounter (Map.unionWith max a b)
+
+gcValue :: GCounter -> Int
+gcValue (GCounter m) = sum (Map.elems m)
+
+gcGet :: String -> GCounter -> Int
+gcGet nodeId (GCounter m) = Map.findWithDefault 0 nodeId m
+
+-- ---------------------------------------------------------------------------
+-- OR-Set: observe-remove set. Each element tagged with a set of UUIDs.
+-- Add: insert a fresh UUID tag. Remove: delete all tags for that element.
+-- Merge: union of tag sets per element. An element is present iff it has
+-- at least one tag.
+-- ---------------------------------------------------------------------------
+
+newtype ORSet a = ORSet (Map a (Set UUID))
+  deriving (Show, Eq, Generic)
+
+instance (Ord a, Binary a) => Binary (ORSet a) where
+  put (ORSet m) = put (Map.toAscList (fmap Set.toAscList m))
+  get = ORSet . Map.fromAscList . map (fmap Set.fromAscList) <$> get
+
+orEmpty :: ORSet a
+orEmpty = ORSet Map.empty
+
+orAdd :: Ord a => UUID -> a -> ORSet a -> ORSet a
+orAdd tag x (ORSet m) =
+  ORSet (Map.insertWith Set.union x (Set.singleton tag) m)
+
+orRemove :: Ord a => a -> ORSet a -> ORSet a
+orRemove x (ORSet m) = ORSet (Map.delete x m)
+
+orMerge :: Ord a => ORSet a -> ORSet a -> ORSet a
+orMerge (ORSet a) (ORSet b) =
+  ORSet (Map.unionWith Set.union a b)
+
+orMembers :: Ord a => ORSet a -> Set a
+orMembers (ORSet m) = Map.keysSet (Map.filter (not . Set.null) m)

src/FL/Data/PaySim.hs

@@ -0,0 +1,231 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE OverloadedStrings #-}
+
+module FL.Data.PaySim
+  ( loadAndPartition
+  -- , syntheticPartition
+  ) where
+
+import Data.ByteString.Lazy qualified as BL
+import Data.Csv qualified as Csv
+import Data.List (sortBy)
+import Data.Ord (comparing)
+import Data.Vector qualified as V
+import FL.NN (LocalData (..))
+import FL.Types (FLConfig (..), NormMode (..), PartitionMode (..))
+import GHC.Generics (Generic)
+import Numeric.LinearAlgebra hiding ((<>))
+import Numeric.LinearAlgebra qualified as LA
+import System.Random (mkStdGen, randomRs)
+
+-- ---------------------------------------------------------------------------
+-- PaySim CSV record
+-- ---------------------------------------------------------------------------
+
+data PaySimRow = PaySimRow
+  { rowStep           :: !Double
+  , rowType           :: !String
+  , rowAmount         :: !Double
+  , rowOldBalOrig     :: !Double
+  , rowNewBalOrig     :: !Double
+  , rowOldBalDest     :: !Double
+  , rowNewBalDest     :: !Double
+  , rowIsFraud        :: !Double
+  } deriving (Show, Generic)
+
+instance Csv.FromNamedRecord PaySimRow where
+  parseNamedRecord r =
+    PaySimRow
+      <$> r Csv..: "step"
+      <*> r Csv..: "type"
+      <*> r Csv..: "amount"
+      <*> r Csv..: "oldbalanceOrg"
+      <*> r Csv..: "newbalanceOrig"
+      <*> r Csv..: "oldbalanceDest"
+      <*> r Csv..: "newbalanceDest"
+      <*> r Csv..: "isFraud"
+
+rowToFeatures :: PaySimRow -> Vector Double
+rowToFeatures PaySimRow{..} =
+  let typeVec = case rowType of
+        "CASH_IN"   -> [1,0,0,0,0]
+        "CASH_OUT"  -> [0,1,0,0,0]
+        "DEBIT"     -> [0,0,1,0,0]
+        "PAYMENT"   -> [0,0,0,1,0]
+        "TRANSFER"  -> [0,0,0,0,1]
+        _           -> [0,0,0,0,0]
+      numeric = [ rowAmount
+                , rowOldBalOrig
+                , rowNewBalOrig
+                , rowOldBalDest
+                , rowNewBalDest
+                , rowNewBalOrig - rowOldBalOrig   -- balance change orig
+                , rowNewBalDest - rowOldBalDest   -- balance change dest
+                , rowAmount / (rowOldBalOrig + 1) -- amount ratio orig
+                , rowAmount / (rowOldBalDest + 1) -- amount ratio dest
+                , rowStep / 744                   -- normalized step (max 744 in PaySim)
+                ]
+      -- pad to 29 features (5 one-hot + 10 numeric + 14 zeros)
+      padded = typeVec ++ numeric ++ replicate 14 0.0
+  in fromList padded
+
+-- ---------------------------------------------------------------------------
+-- Load PaySim CSV
+-- ---------------------------------------------------------------------------
+
+loadPaySimCSV :: FilePath -> IO (Either String (Matrix Double, Vector Double))
+loadPaySimCSV path = do
+  raw <- BL.readFile path
+  case Csv.decodeByName raw of
+    Left err -> return (Left err)
+    Right (_, rows) -> do
+      let featureList = V.toList (V.map rowToFeatures rows)
+          labelList   = V.toList (V.map rowIsFraud rows)
+      return (Right (fromRows featureList, fromList labelList))
+
+-- ---------------------------------------------------------------------------
+-- Normalization
+-- ---------------------------------------------------------------------------
+
+normalizeMinMax :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
+normalizeMinMax m =
+  let mins  = fromList [minElement (m ¿ [j]) | j <- [0 .. cols m - 1]]
+      maxs  = fromList [maxElement (m ¿ [j]) | j <- [0 .. cols m - 1]]
+      range = cmap (\x -> if x == 0 then 1 else x) (maxs - mins)
+      norm  = (m - LA.repmat (asRow mins) (rows m) 1)
+              / LA.repmat (asRow range)  (rows m) 1
+  in (norm, mins, maxs)
+
+normalizeZScore :: Matrix Double -> (Matrix Double, Vector Double, Vector Double)
+normalizeZScore m =
+  let n     = fromIntegral (rows m) :: Double
+      means = fromList [sumElements (m ¿ [j]) / n | j <- [0 .. cols m - 1]]
+      vars  = fromList [sumElements (cmap (^(2::Int)) ((m ¿ [j]) - scalar (means ! j))) / n
+                       | j <- [0 .. cols m - 1]]
+      stds  = cmap (\v -> if v == 0 then 1 else sqrt v) vars
+      norm  = (m - LA.repmat (asRow means) (rows m) 1)
+              / LA.repmat (asRow stds)  (rows m) 1
+  in (norm, means, stds)
+
+applyNorm :: NormMode -> (Matrix Double, Vector Double, Vector Double) -> Matrix Double -> Matrix Double
+applyNorm mode (_, p1, p2) m =
+  case mode of
+    MinMax ->
+      let range = cmap (\x -> if x == 0 then 1 else x) (p2 - p1)
+      in (m - LA.repmat (asRow p1) (rows m) 1) / LA.repmat (asRow range) (rows m) 1
+    ZScore ->
+      let stds = cmap (\x -> if x == 0 then 1 else x) p2
+      in (m - LA.repmat (asRow p1) (rows m) 1) / LA.repmat (asRow stds)  (rows m) 1
+
+-- ---------------------------------------------------------------------------
+-- IID partition: deterministic shard by trainer index
+-- ---------------------------------------------------------------------------
+
+partitionIID :: Int -> Int -> Matrix Double -> Vector Double -> LocalData
+partitionIID trainerIdx numTrainers feats labels =
+  let n       = rows feats
+      -- Shuffle with fixed seed for reproducibility across processes
+      seed    = 2024
+      rng     = mkStdGen seed
+      sorted  = sortBy (comparing fst)
+                  (zip (randomRs (0.0 :: Double, 1.0) rng) [0 .. n - 1])
+      idxs    = map snd sorted
+      shard   = chunkedShard trainerIdx numTrainers idxs
+      f       = feats ? shard
+      l       = fromList [labels ! i | i <- shard]
+  in LocalData f l
+
+chunkedShard :: Int -> Int -> [a] -> [a]
+chunkedShard idx total xs =
+  let n       = length xs
+      size    = n `div` total
+      start   = idx * size
+      end     = if idx == total - 1 then n else start + size
+  in take (end - start) (drop start xs)
+
+-- ---------------------------------------------------------------------------
+-- Non-IID partition: each trainer gets dominant transaction type
+-- (approximated via deterministic row assignment)
+-- ---------------------------------------------------------------------------
+
+partitionNonIID :: Int -> Int -> Matrix Double -> Vector Double -> LocalData
+partitionNonIID trainerIdx numTrainers feats labels =
+  let n         = rows feats
+      -- Assign rows to trainers based on dominant type one-hot (cols 0–4)
+      -- Trainer i "owns" type (i mod 5); remaining trainers share type 0
+      ownType   = trainerIdx `mod` 5
+      isOwn i   = let row = toList (flatten (feats ? [i]))
+                  in length row > ownType && row !! ownType > 0.5
+      -- With more than 5 trainers several trainers own the same type;
+      -- split that type's rows among them instead of giving each trainer
+      -- all of them (which duplicated data across shards).
+      sharers   = [i | i <- [0 .. numTrainers - 1], i `mod` 5 == ownType]
+      shareIdx  = length (takeWhile (/= trainerIdx) sharers)
+      -- Gather own-type rows + 10% of other rows for minority
+      ownIdxs   = chunkedShard shareIdx (length sharers) [i | i <- [0 .. n-1], isOwn i]
+      otherIdxs = [i | i <- [0 .. n-1], not (isOwn i)]
+      minority  = chunkedShard trainerIdx numTrainers otherIdxs
+      allIdxs   = ownIdxs ++ minority
+      f         = feats ? allIdxs
+      l         = fromList [labels ! i | i <- allIdxs]
+  in LocalData f l
+
+-- ---------------------------------------------------------------------------
+-- Synthetic data fallback
+-- ---------------------------------------------------------------------------
+
+syntheticData :: Int -> Int -> Int -> LocalData
+syntheticData trainerIdx numTrainers nSamples =
+  let seed    = trainerIdx * 13337 + numTrainers
+      rng     = mkStdGen seed
+      -- 29 features from normal distribution
+      featVals = take (nSamples * 29) (randomRs (-3.0 :: Double, 3.0) rng)
+      feats   = (nSamples >< 29) featVals
+      -- Base 2% fraud rate, +1% per trainer index for variety (real PaySim
+      -- is ~0.1%, but synthetic runs need enough positives to learn from)
+      fraudRate = 0.02 + fromIntegral trainerIdx * 0.01
+      rng2    = mkStdGen (seed + 1)
+      labelVals = take nSamples (randomRs (0.0 :: Double, 1.0) rng2)
+      labels  = fromList [if v < fraudRate then 1.0 else 0.0 | v <- labelVals]
+  in LocalData feats labels
+
+-- ---------------------------------------------------------------------------
+-- Top-level entry point used by Provider and P2P startup
+-- ---------------------------------------------------------------------------
+
+-- Returns (train partition, test data).
+-- Test data is the first 10% of the full dataset (at least 1000 rows),
+-- taken before partitioning so every node evaluates on the same split.
+loadAndPartition
+  :: FLConfig
+  -> Int    -- this node's trainer index (0-based)
+  -> IO (LocalData, LocalData)
+loadAndPartition FLConfig{..} trainerIdx = do
+  result <- case cfgDataPath of
+    Nothing   -> return (Left "no data path")
+    Just path -> loadPaySimCSV path
+  case result of
+    Left _ -> do
+      -- Synthetic fallback
+      let testData  = syntheticData (cfgNumTrainers + 1) cfgNumTrainers 2000
+          trainData = syntheticData trainerIdx cfgNumTrainers 10000
+      return (trainData, testData)
+    Right (feats, labels) -> do
+      let n         = rows feats
+          -- Aim for 1000 test rows for stable metrics, but never take more
+          -- than half the data — max alone made testN exceed n on small
+          -- CSVs and crashed the train subMatrix with a negative size.
+          testN     = min (max 1000 (n `div` 10)) (n `div` 2)
+          testFeats = subMatrix (0, 0) (testN, cols feats) feats
+          testLabs  = subVector 0 testN labels
+          trainF    = subMatrix (testN, 0) (n - testN, cols feats) feats
+          trainL    = subVector testN (n - testN) labels
+          (normF, p1, p2) = case cfgNorm of
+            MinMax -> normalizeMinMax trainF
+            ZScore -> normalizeZScore trainF
+          normTest  = applyNorm cfgNorm (normF, p1, p2) testFeats
+          trainData = case cfgPartition of
+            IID    -> partitionIID    trainerIdx cfgNumTrainers normF trainL
+            NonIID -> partitionNonIID trainerIdx cfgNumTrainers normF trainL
+          testData  = LocalData normTest testLabs
+      return (trainData, testData)

src/FL/NN.hs

@@ -0,0 +1,290 @@
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE StrictData #-}
+
+module FL.NN
+  ( MLPWeights (..)
+  , LocalData (..)
+  , initWeights
+  , trainLocal
+  , evaluate
+  , flattenWeights
+  , unflattenWeights
+  , fedAvg
+  , localDataSize
+  ) where
+
+import Data.List (foldl', sortBy)
+import Data.Ord (Down (..), comparing)
+import FL.Types (FLConfig (..), Metrics (..))
+import Numeric.LinearAlgebra hiding ((<>))
+import Numeric.LinearAlgebra qualified as LA
+import System.Random (mkStdGen, randomRs)
+
+-- ---------------------------------------------------------------------------
+-- Types
+-- ---------------------------------------------------------------------------
+
+data MLPWeights = MLPWeights
+  { w1 :: Matrix Double  -- 64 × 29
+  , b1 :: Vector Double  -- 64
+  , w2 :: Matrix Double  -- 32 × 64
+  , b2 :: Vector Double  -- 32
+  , w3 :: Matrix Double  --  2 × 32
+  , b3 :: Vector Double  --  2
+  }
+
+data LocalData = LocalData
+  { ldFeatures :: Matrix Double  -- N × 29
+  , ldLabels   :: Vector Double  -- N  (0.0 or 1.0)
+  }
+
+localDataSize :: LocalData -> Int
+localDataSize = rows . ldFeatures
+
+-- ---------------------------------------------------------------------------
+-- Weight initialization (Xavier uniform)
+-- ---------------------------------------------------------------------------
+
+xavierMatrix :: Int -> Int -> Int -> Matrix Double
+xavierMatrix seed fanIn fanOut =
+  let bound = sqrt (6.0 / fromIntegral (fanIn + fanOut))
+      n     = fanIn * fanOut
+      vals  = take n (randomRs (-bound, bound) (mkStdGen seed))
+  in (fanOut >< fanIn) vals
+
+initWeights :: IO MLPWeights
+initWeights = do
+  -- Use fixed seeds so multiple processes with same call order get different
+  -- weight initializations only if we pass different seeds; for now deterministic.
+  let w1' = xavierMatrix 42  29 64
+      b1' = konst 0 64
+      w2' = xavierMatrix 137 64 32
+      b2' = konst 0 32
+      w3' = xavierMatrix 999 32  2
+      b3' = konst 0 2
+  return MLPWeights{ w1=w1', b1=b1', w2=w2', b2=b2', w3=w3', b3=b3' }
+
+-- ---------------------------------------------------------------------------
+-- Activation functions
+-- ---------------------------------------------------------------------------
+
+relu :: Matrix Double -> Matrix Double
+relu = cmap (max 0)
+
+reluGrad :: Matrix Double -> Matrix Double
+reluGrad = cmap (\x -> if x > 0 then 1 else 0)
+
+softmaxRows :: Matrix Double -> Matrix Double
+softmaxRows m = fromRows (map softmaxVec (toRows m))
+  where
+    -- Subtract the row max before exponentiating so large logits
+    -- can't overflow exp and poison the weights with NaN.
+    softmaxVec v =
+      let mx = maxElement v
+          e  = cmap (\x -> exp (x - mx)) v
+          s  = sumElements e
+      in cmap (/ s) e
+
+-- ---------------------------------------------------------------------------
+-- Forward pass
+-- ---------------------------------------------------------------------------
+
+-- Returns (output probs N×2, z1, a1, z2, a2) for backprop.
+forwardWithCache
+  :: MLPWeights
+  -> Matrix Double  -- N × 29
+  -> (Matrix Double, Matrix Double, Matrix Double, Matrix Double, Matrix Double)
+forwardWithCache MLPWeights{..} x =
+  let z1  = x LA.<> tr w1 + repRows (rows x) b1
+      a1  = relu z1
+      z2  = a1 LA.<> tr w2 + repRows (rows x) b2
+      a2  = relu z2
+      z3  = a2 LA.<> tr w3 + repRows (rows x) b3
+      out = softmaxRows z3
+  in (out, z1, a1, z2, a2)
+
+repRows :: Int -> Vector Double -> Matrix Double
+repRows n v = fromRows (replicate n v)
+
+colSums :: Matrix Double -> Vector Double
+colSums m = fromList [sumElements (m ¿ [j]) | j <- [0 .. cols m - 1]]
+
+-- ---------------------------------------------------------------------------
+-- Loss
+-- ---------------------------------------------------------------------------
+
+crossEntropyLoss :: Matrix Double -> Vector Double -> Double
+crossEntropyLoss probs labels =
+  let n = fromIntegral (rows probs)
+      logProbs = cmap (log . max 1e-15) probs
+      oneHot   = buildOneHot (rows probs) (map round (toList labels))
+      loss     = negate (sumElements (oneHot * logProbs)) / n
+  in loss
+
+buildOneHot :: Int -> [Int] -> Matrix Double
+buildOneHot n ys =
+  fromRows [fromList [if i == y then 1.0 else 0.0 | i <- [0, 1]] | y <- ys]
+
+-- ---------------------------------------------------------------------------
+-- Backward pass (mini-batch SGD)
+-- ---------------------------------------------------------------------------
+
+stepBatch
+  :: MLPWeights
+  -> Matrix Double  -- batch features N × 29
+  -> Vector Double  -- batch labels N
+  -> Double         -- learning rate
+  -> (MLPWeights, Double)
+stepBatch ws@MLPWeights{..} x labels lr =
+  let n              = fromIntegral (rows x) :: Double
+      (out, z1, a1, z2, a2) = forwardWithCache ws x
+      oneHot         = buildOneHot (rows x) (map round (toList labels))
+      -- Output delta: dL/dZ3 = (probs - oneHot) / N
+      dZ3            = cmap (/ n) (out - oneHot)           -- N × 2
+      dW3            = tr dZ3 LA.<> a2                     -- 2 × 32
+      db3            = colSums dZ3                         -- 2
+      dA2            = dZ3 LA.<> w3                        -- N × 32
+      dZ2            = dA2 * reluGrad z2                   -- N × 32
+      dW2            = tr dZ2 LA.<> a1                     -- 32 × 64
+      db2            = colSums dZ2                         -- 32
+      dA1            = dZ2 LA.<> w2                        -- N × 64
+      dZ1            = dA1 * reluGrad z1                   -- N × 64
+      dW1            = tr dZ1 LA.<> x                      -- 64 × 29
+      db1            = colSums dZ1                         -- 64
+      ws'            = MLPWeights
+        { w1 = w1 - cmap (* lr) dW1
+        , b1 = b1 - cmap (* lr) db1
+        , w2 = w2 - cmap (* lr) dW2
+        , b2 = b2 - cmap (* lr) db2
+        , w3 = w3 - cmap (* lr) dW3
+        , b3 = b3 - cmap (* lr) db3
+        }
+      loss = crossEntropyLoss out labels
+  in (ws', loss)
+
+-- ---------------------------------------------------------------------------
+-- Training
+-- ---------------------------------------------------------------------------
+
+-- Shuffle row indices with a seed, then split into mini-batches.
+shuffleIndices :: Int -> Int -> [Int]
+shuffleIndices seed n =
+  let rng = mkStdGen seed
+      tagged = zip (randomRs (0.0 :: Double, 1.0) rng) [0 .. n - 1]
+  in map snd (sortBy (comparing fst) tagged)
+
+trainEpoch
+  :: MLPWeights
+  -> Matrix Double  -- features N × 29
+  -> Vector Double  -- labels N
+  -> Int            -- batch size
+  -> Double         -- learning rate
+  -> Int            -- seed for shuffle
+  -> (MLPWeights, Double)
+trainEpoch ws feats labels batchSize lr seed =
+  let n       = rows feats
+      idxs    = shuffleIndices seed n
+      batches = chunksOf batchSize idxs
+      step (w, totalLoss) batch =
+        let bFeats  = feats ? batch
+            bLabels = fromList [labels ! i | i <- batch]
+            (w', l) = stepBatch w bFeats bLabels lr
+        in (w', totalLoss + l)
+      (wFinal, lossSum) = foldl' step (ws, 0.0) batches
+  in (wFinal, lossSum / fromIntegral (length batches))
+
+chunksOf :: Int -> [a] -> [[a]]
+chunksOf _ [] = []
+chunksOf n xs = let (h, t) = splitAt n xs in h : chunksOf n t
+
+-- | Train for the configured number of epochs. The shuffle seed is derived
+-- from the trainer id and round so batches differ across rounds and trainers
+-- (a bare epoch counter would replay the identical shuffle every round).
+trainLocal :: FLConfig -> String -> Int -> MLPWeights -> LocalData -> IO MLPWeights
+trainLocal FLConfig{..} trainerId roundNum initW LocalData{..} = do
+  let seedBase = foldl' (\h c -> 31 * h + fromEnum c) (roundNum * 7919) trainerId
+      go ws epoch =
+        let (ws', _loss) = trainEpoch ws ldFeatures ldLabels cfgBatchSize cfgLearningRate (seedBase + epoch)
+        in if epoch >= cfgEpochs then ws'
+           else go ws' (epoch + 1)
+  return (go initW 1)
+
+-- ---------------------------------------------------------------------------
+-- Evaluation
+-- ---------------------------------------------------------------------------
+
+evaluate :: MLPWeights -> LocalData -> Int -> Metrics
+evaluate ws LocalData{..} round_ =
+  let n         = rows ldFeatures
+      (probs, _, _, _, _) = forwardWithCache ws ldFeatures
+      predLabels = fromList [if (probs ! i ! 1) >= 0.5 then 1.0 else 0.0 | i <- [0 .. n - 1]]
+      scores     = fromList [probs ! i ! 1 | i <- [0 .. n - 1]]
+      tp = sumElements (predLabels * ldLabels)
+      fp = sumElements (predLabels * (1 - ldLabels))
+      fn = sumElements ((1 - predLabels) * ldLabels)
+      tn = sumElements ((1 - predLabels) * (1 - ldLabels))
+      acc  = (tp + tn) / fromIntegral n
+      prec = if tp + fp == 0 then 0 else tp / (tp + fp)
+      rec  = if tp + fn == 0 then 0 else tp / (tp + fn)
+      f1   = if prec + rec == 0 then 0 else 2 * prec * rec / (prec + rec)
+      auc  = computeAUCROC (toList scores) (toList ldLabels)
+  in Metrics acc prec rec f1 auc round_
+
+computeAUCROC :: [Double] -> [Double] -> Double
+computeAUCROC scores labels =
+  let nPos = fromIntegral (length (filter (== 1) labels))
+      nNeg = fromIntegral (length (filter (== 0) labels))
+  in if nPos == 0 || nNeg == 0
+     then 0.5
+     else
+       let sorted = sortBy (comparing (Down . fst)) (zip scores labels)
+           points = scanl (\(tp, fp) (_, y) -> (tp + y, fp + (1 - y))) (0, 0) sorted
+           trapz (tp0, fp0) (tp1, fp1) =
+             let fpr0 = fp0 / nNeg; tpr0 = tp0 / nPos
+                 fpr1 = fp1 / nNeg; tpr1 = tp1 / nPos
+             in (fpr1 - fpr0) * (tpr0 + tpr1) / 2.0
+       in sum (zipWith trapz points (tail points))
+
+-- ---------------------------------------------------------------------------
+-- Serialization bridge
+-- ---------------------------------------------------------------------------
+
+flattenWeights :: MLPWeights -> [Double]
+flattenWeights MLPWeights{..} =
+  toList (flatten w1)
+  ++ toList b1
+  ++ toList (flatten w2)
+  ++ toList b2
+  ++ toList (flatten w3)
+  ++ toList b3
+
+-- Fixed architecture: 29→64→32→2
+-- w1: 64*29=1856, b1: 64, w2: 32*64=2048, b2: 32, w3: 2*32=64, b3: 2
+unflattenWeights :: [Double] -> MLPWeights
+unflattenWeights xs =
+  let (xw1, r1) = splitAt (64*29) xs
+      (xb1, r2) = splitAt 64     r1
+      (xw2, r3) = splitAt (32*64) r2
+      (xb2, r4) = splitAt 32     r3
+      (xw3, r5) = splitAt (2*32)  r4
+      xb3        = take 2 r5
+  in MLPWeights
+       { w1 = (64 >< 29) xw1
+       , b1 = fromList xb1
+       , w2 = (32 >< 64) xw2
+       , b2 = fromList xb2
+       , w3 = ( 2 >< 32) xw3
+       , b3 = fromList xb3
+       }
+
+-- ---------------------------------------------------------------------------
+-- FedAvg
+-- ---------------------------------------------------------------------------
+
+-- Weighted average of flat weight vectors.
+fedAvg :: [(Int, [Double])] -> [Double]
+fedAvg [] = []
+fedAvg updates =
+  let totalN   = fromIntegral (sum (map fst updates)) :: Double
+      weighted (n, ws) = map (* (fromIntegral n / totalN)) ws
+  in foldl1 (zipWith (+)) (map weighted updates)

src/FL/P2P.hs

@@ -0,0 +1,182 @@
+{-# LANGUAGE RecordWildCards #-}
+module FL.P2P
+  ( runPeer
+  ) where
+
+import Control.Actor
+import Control.Actor.Registry (registerActor, lookupRemoteActor)
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.MVar (newEmptyMVar, takeMVar)
+import Control.Concurrent.STM
+  ( TMVar, atomically, newEmptyTMVarIO, putTMVar, readTMVar, tryTakeTMVar )
+import Control.Monad (forM, forM_, void)
+import Control.Monad.IO.Class (liftIO)
+import FL.Actors.Logger (logFn, logFn', loggerHandler)
+import FL.Actors.PeerCoordinator
+import FL.CRDT (gcEmpty, orEmpty)
+import FL.Data.PaySim (loadAndPartition)
+import FL.NN
+  ( LocalData, evaluate, flattenWeights, initWeights
+  , localDataSize, trainLocal, unflattenWeights
+  )
+import FL.Retry (retryConnect, retryLookup)
+import FL.Types
+
+-- ---------------------------------------------------------------------------
+-- P2P-specific actor state types
+-- ---------------------------------------------------------------------------
+
+-- Evaluator reads the PeerCoordinator ref lazily from a TMVar.
+-- This breaks the circular dependency: evaluator is spawned before PC,
+-- but PC's ref is filled into the TMVar right after PC spawns.
+data P2PEvalState = P2PEvalState
+  { p2eTestData :: LocalData
+  , p2ePcCell   :: TMVar (ActorRef PeerCoordinatorMsg ())
+  , p2eLogRef   :: ActorRef LogMsg ()
+  }
+
+p2pEvalHandler :: Handler EvaluatorMsg P2PEvalState ()
+p2pEvalHandler (EvaluateModel gm) = do
+  st <- state
+  let ws      = unflattenWeights (wData (gmWeights gm))
+      metrics = evaluate ws (p2eTestData st) (gmRound gm)
+  logFn (p2eLogRef st) "Evaluator" INFO $
+    "P2P round " <> show (gmRound gm)
+    <> " acc=" <> show4 (mAccuracy metrics)
+    <> " auc=" <> show4 (mAUCROC  metrics)
+  pcRef <- liftIO $ atomically $ readTMVar (p2ePcCell st)
+  cast (PeerEvaluationDone metrics) pcRef
+  pass
+  where
+    show4 x = show (fromIntegral (round (x * 10000) :: Int) / 10000.0 :: Double)
+
+-- Trainer likewise reads the PeerCoordinator ref lazily.
+data P2PTrainerState = P2PTrainerState
+  { p2tId        :: String
+  , p2tData      :: LocalData
+  , p2tLastRound :: Int
+  , p2tPcCell    :: TMVar (ActorRef PeerCoordinatorMsg ())
+  , p2tLogRef    :: ActorRef LogMsg ()
+  , p2tConfig    :: FLConfig
+  }
+
+p2pTrainerHandler :: Handler TrainerMsg P2PTrainerState ()
+p2pTrainerHandler (StartTraining gm) = do
+  st <- state
+  let round_ = gmRound gm
+  if round_ <= p2tLastRound st
+    then pass
+    else do
+      logFn (p2tLogRef st) (p2tId st) INFO $
+        "P2P training round " <> show round_
+      let initW = unflattenWeights (wData (gmWeights gm))
+      updatedW <- liftIO $ trainLocal (p2tConfig st) (p2tId st) round_ initW (p2tData st)
+      let n   = localDataSize (p2tData st)
+          upd = ModelUpdate (p2tId st) round_ n (Weights (flattenWeights updatedW))
+      pcRef <- liftIO $ atomically $ readTMVar (p2tPcCell st)
+      cast (PeerTrainingDone upd) pcRef
+      passWith st { p2tLastRound = round_ }
+p2pTrainerHandler TrainerStop = pass
+
+-- ---------------------------------------------------------------------------
+-- Peer node startup
+-- ---------------------------------------------------------------------------
+
+runPeer :: FLConfig -> Int -> [NodeAddr] -> IO ()
+runPeer cfg peerIdx knownPeers = do
+  let FLConfig{..} = cfg
+      peerPort = cfgBasePort
+      peerId   = "peer-" <> show peerIdx
+
+  rt      <- initRuntime (NodeAddr "localhost" (fromIntegral peerPort)) createTCPTransport
+  doneVar <- newEmptyMVar
+
+  withRuntime rt $ do
+    forM_ knownPeers $ \addr -> retryConnect 20 addr
+    liftIO $ threadDelay 300_000
+
+    logRef <- spawnActor loggerHandler stopOnDeath ()
+    registerActor (peerId <> "-logger") logRef
+
+    initW <- liftIO initWeights
+    let initModel = GlobalModel 0 (Weights (flattenWeights initW))
+
+    (trainData, testData) <- liftIO $ loadAndPartition cfg peerIdx
+
+    -- Shared TMVar: filled with the PeerCoordinator ref after it spawns.
+    -- Evaluator and Trainer read from it lazily on their first message.
+    pcCell <- liftIO newEmptyTMVarIO
+
+    evalRef <- spawnActor p2pEvalHandler stopOnDeath
+      (P2PEvalState testData pcCell logRef)
+
+    trainerRef <- spawnActor p2pTrainerHandler stopOnDeath
+      P2PTrainerState
+        { p2tId        = peerId
+        , p2tData      = trainData
+        , p2tLastRound = -1
+        , p2tPcCell    = pcCell
+        , p2tLogRef    = logRef
+        , p2tConfig    = cfg
+        }
+
+    -- PeerCoordinator — placeholder self-ref (replaced by SetPeers)
+    let dummyId  = ActorId 0 (read "00000000-0000-0000-0000-000000000002")
+        dummySelf = RemoteRef dummyId :: ActorRef PeerCoordinatorMsg ()
+
+    let pcSt = PeerCoordinatorState
+          { pcsNodeId       = peerId
+          , pcsRoundCounter = gcEmpty
+          , pcsParticipants = orEmpty
+          , pcsCurrentRound = 0
+          , pcsTotalRounds  = cfgRounds
+          , pcsGlobalModel  = initModel
+          , pcsPeerRefs     = []
+          , pcsSelfRef      = dummySelf
+          , pcsTrainerRef   = trainerRef
+          , pcsEvalRef      = evalRef
+          , pcsLogRef       = logRef
+          , pcsCollected    = []
+          , pcsNumPeers     = cfgNumTrainers
+          , pcsDoneVar      = doneVar
+          }
+
+    -- Supervised PeerCoordinator. On every (re)spawn the shared cell is
+    -- refreshed and the registry name re-registered, so trainer/evaluator
+    -- keep reaching the live incarnation. (State restarts fresh.)
+    supervise' OneForOne
+      [ ChildSpec
+          { csRun = \target -> do
+              ref <- spawnActor peerCoordinatorHandler peerCoordinatorDeathHandler pcSt
+              linkActorTo target ref
+              return ref
+          , csOnSpawn = \ref -> do
+              liftIO $ atomically $ do
+                void $ tryTakeTMVar pcCell
+                putTMVar pcCell ref
+              registerActor ("peercoord-" <> peerId) ref
+          }
+      ]
+    pcRef <- liftIO $ atomically $ readTMVar pcCell
+
+    logFn' logRef peerId INFO $
+      "Peer ready, waiting for " <> show (cfgNumTrainers - 1) <> " other peers"
+
+    -- Discover all other peer coordinators
+    let otherIds = ["peer-" <> show i | i <- [0 .. cfgNumTrainers - 1], i /= peerIdx]
+    peerPcRefs <- forM otherIds $ \pid ->
+      retryLookup 40
+        (lookupRemoteActor ("peercoord-" <> pid)
+          :: RuntimeM (Maybe (ActorRef PeerCoordinatorMsg ())))
+
+    -- Wire up PeerCoordinator with peers
+    let peerIds = map actorRefId peerPcRefs
+    cast' (SetPeers peerIds) pcRef
+
+    liftIO $ threadDelay 200_000
+    cast' PeerStartRound pcRef
+    logFn' logRef peerId INFO "P2P round 1 started"
+
+  takeMVar doneVar
+  threadDelay 300_000
+  putStrLn $ "\n=== Peer " <> show peerIdx <> " federation complete ==="

src/FL/Provider.hs

@@ -0,0 +1,149 @@
+{-# LANGUAGE RecordWildCards #-}
+module FL.Provider
+  ( runCoordinator
+  , runTrainer
+  ) where
+
+import Control.Actor
+import Control.Actor.Registry (lookupRemoteActor, registerActor)
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.MVar (newEmptyMVar, takeMVar)
+import Control.Concurrent.STM
+  ( atomically, newEmptyTMVarIO, putTMVar, readTMVar, tryTakeTMVar )
+import Control.Monad (void)
+import Control.Monad.IO.Class (liftIO)
+import Data.Map.Strict qualified as Map
+import FL.Actors.Aggregator
+import FL.Actors.Coordinator
+import FL.Actors.Evaluator
+import FL.Actors.Logger
+import FL.Actors.Trainer
+import FL.CRDT (orEmpty)
+import FL.Data.PaySim (loadAndPartition)
+import FL.NN (flattenWeights, initWeights)
+import FL.Retry (retryConnect, retryLookup)
+import FL.Types
+
+-- ---------------------------------------------------------------------------
+-- Coordinator node
+-- ---------------------------------------------------------------------------
+
+runCoordinator :: FLConfig -> IO ()
+runCoordinator cfg = do
+  let FLConfig{..} = cfg
+  rt      <- initRuntime (NodeAddr "localhost" (fromIntegral cfgBasePort)) createTCPTransport
+  doneVar <- newEmptyMVar
+
+  withRuntime rt $ do
+    -- Logger
+    logRef <- spawnActor loggerHandler stopOnDeath ()
+    registerActor "logger" logRef
+
+    -- Initial global model
+    initW <- liftIO initWeights
+    let initModel = GlobalModel 0 (Weights (flattenWeights initW))
+
+    -- Shared TMVar for coordinator ref: filled after coordinator spawns.
+    -- Aggregator and Evaluator read from it lazily.
+    coordCell <- liftIO newEmptyTMVarIO
+
+    -- Load test data (coordinator holds the shared test split)
+    (_, testData) <- liftIO $ loadAndPartition cfg 0
+
+    -- Evaluator
+    evalRef <- spawnActor evaluatorHandler stopOnDeath
+      (EvaluatorState testData coordCell logRef)
+    registerActor "evaluator" evalRef
+
+    -- Aggregator
+    aggRef <- spawnActor aggregatorHandler stopOnDeath
+      (AggregatorState cfgNumTrainers [] coordCell logRef 0)
+    registerActor "aggregator" aggRef
+
+    -- Coordinator (knows aggRef and evalRef at spawn time)
+    let initCoordSt = CoordinatorState
+          { csTrainers         = orEmpty
+          , csTrainerRefs      = Map.empty
+          , csExpectedTrainers = cfgNumTrainers
+          , csCurrentRound     = 0
+          , csTotalRounds      = cfgRounds
+          , csGlobalModel      = initModel
+          , csAggRef           = aggRef
+          , csEvalRef          = evalRef
+          , csLogRef           = logRef
+          , csDoneVar          = doneVar
+          }
+    -- Supervised coordinator. On every (re)spawn — including restarts — the
+    -- shared cell is refreshed and the registry name re-registered, so
+    -- aggregator/evaluator/trainers keep reaching the live incarnation.
+    -- (Coordinator state itself restarts fresh: round progress is lost.)
+    supervise' OneForOne
+      [ ChildSpec
+          { csRun = \target -> do
+              ref <- spawnActor coordinatorHandler coordinatorDeathHandler initCoordSt
+              linkActorTo target ref
+              return ref
+          , csOnSpawn = \ref -> do
+              liftIO $ atomically $ do
+                void $ tryTakeTMVar coordCell
+                putTMVar coordCell ref
+              registerActor "coordinator" ref
+          }
+      ]
+    void $ liftIO $ atomically $ readTMVar coordCell
+
+    logFn' logRef "System" INFO $
+      "Coordinator ready, waiting for " <> show cfgNumTrainers <> " trainers"
+
+  takeMVar doneVar
+  threadDelay 300_000
+  putStrLn "\n=== Provider federation complete ==="
+
+-- ---------------------------------------------------------------------------
+-- Trainer node
+-- ---------------------------------------------------------------------------
+
+runTrainer :: FLConfig -> Int -> NodeAddr -> IO ()
+runTrainer cfg trainerIdx coordAddr = do
+  let FLConfig{..} = cfg
+  let trainerPort = cfgBasePort + trainerIdx + 1
+      trainerId   = "trainer-" <> show trainerIdx
+
+  rt <- initRuntime (NodeAddr "localhost" (fromIntegral trainerPort)) createTCPTransport
+
+  withRuntime rt $ do
+    retryConnect 20 coordAddr
+    liftIO $ threadDelay 500_000  -- wait for registry sync
+
+    -- Look up aggregator and logger on the coordinator node (retry until
+    -- available; the logger is registered first, so if the aggregator
+    -- resolves the logger will too)
+    aggRef <- retryLookup 20
+                (lookupRemoteActor "aggregator" :: RuntimeM (Maybe (ActorRef AggregatorMsg ())))
+    logRef <- retryLookup 20
+                (lookupRemoteActor "logger" :: RuntimeM (Maybe (ActorRef LogMsg ())))
+
+    (trainData, _) <- liftIO $ loadAndPartition cfg trainerIdx
+
+    let trainerSt = TrainerState
+          { tsId        = trainerId
+          , tsData      = trainData
+          , tsLastRound = -1
+          , tsAggRef    = aggRef
+          , tsLogRef    = logRef
+          , tsConfig    = cfg
+          }
+    trainerRef <- spawnActor trainerHandler stopOnDeath trainerSt
+    registerActor trainerId trainerRef
+
+    -- Register with coordinator
+    coordRef <- retryLookup 20
+                  (lookupRemoteActor "coordinator" :: RuntimeM (Maybe (ActorRef CoordinatorMsg ())))
+    cast' (RegisterTrainer trainerId (actorRefId trainerRef)) coordRef
+
+    logFn' logRef trainerId INFO "Registered with coordinator"
+
+  waitForever
+
+waitForever :: IO ()
+waitForever = threadDelay maxBound >> waitForever

src/FL/Retry.hs

@@ -0,0 +1,34 @@
+-- | Startup helpers shared by the provider and P2P entry points: retry a
+-- node connection or a registry lookup until it succeeds or attempts run out.
+module FL.Retry
+  ( retryConnect
+  , retryLookup
+  ) where
+
+import Control.Actor (ActorRef, NodeAddr, RuntimeM, connect, withRuntime)
+import Control.Concurrent (threadDelay)
+import Control.Exception (SomeException, try)
+import Control.Monad (void)
+import Control.Monad.IO.Class (liftIO)
+import Control.Monad.Reader (ask)
+
+retryConnect :: Int -> NodeAddr -> RuntimeM ()
+retryConnect 0 addr = error $ "retryConnect: could not connect to " <> show addr
+retryConnect n addr = do
+  rt     <- ask
+  result <- liftIO $ try @SomeException $ withRuntime rt $ void (connect Nothing addr)
+  case result of
+    Right () -> return ()
+    Left  _  -> do
+      liftIO $ threadDelay 1_000_000
+      retryConnect (n - 1) addr
+
+retryLookup :: Int -> RuntimeM (Maybe (ActorRef msg r)) -> RuntimeM (ActorRef msg r)
+retryLookup 0 _ = error "retryLookup: exhausted retries"
+retryLookup n action = do
+  result <- action
+  case result of
+    Just ref -> return ref
+    Nothing  -> do
+      liftIO $ threadDelay 500_000
+      retryLookup (n - 1) action

src/FL/Types.hs

@@ -0,0 +1,160 @@
+module FL.Types
+  ( -- Config
+    FLConfig (..)
+  , PartitionMode (..)
+  , NormMode (..)
+    -- Wire types
+  , Weights (..)
+  , GlobalModel (..)
+  , ModelUpdate (..)
+  , Metrics (..)
+    -- Logger messages
+  , LogLevel (..)
+  , LogEntry (..)
+  , LogMsg (..)
+    -- Actor messages
+  , EvaluatorMsg (..)
+  , TrainerMsg (..)
+  , AggregatorMsg (..)
+  , CoordinatorMsg (..)
+  , PeerSyncPayload (..)
+  , PeerCoordinatorMsg (..)
+  ) where
+
+import Control.Actor.Types (ActorId)
+import Data.Binary (Binary)
+import FL.CRDT (GCounter, ORSet)
+import GHC.Generics (Generic)
+
+-- ---------------------------------------------------------------------------
+-- Config (not a wire type — parsed from CLI, passed around locally)
+-- ---------------------------------------------------------------------------
+
+data PartitionMode = IID | NonIID
+  deriving (Show, Eq, Generic)
+instance Binary PartitionMode
+
+data NormMode = MinMax | ZScore
+  deriving (Show, Eq, Generic)
+instance Binary NormMode
+
+data FLConfig = FLConfig
+  { cfgRounds       :: Int
+  , cfgEpochs       :: Int
+  , cfgLearningRate :: Double
+  , cfgBatchSize    :: Int
+  , cfgNumTrainers  :: Int
+  , cfgDataPath     :: Maybe FilePath
+  , cfgPartition    :: PartitionMode
+  , cfgNorm         :: NormMode
+  , cfgBasePort     :: Int
+  } deriving (Show, Generic)
+
+-- ---------------------------------------------------------------------------
+-- Wire types (sent over the network, must be Binary)
+-- ---------------------------------------------------------------------------
+
+newtype Weights = Weights { wData :: [Double] }
+  deriving (Show, Eq, Generic)
+instance Binary Weights
+
+data GlobalModel = GlobalModel
+  { gmRound   :: Int
+  , gmWeights :: Weights
+  } deriving (Show, Eq, Generic)
+instance Binary GlobalModel
+
+data ModelUpdate = ModelUpdate
+  { muTrainerId  :: String
+  , muRound      :: Int
+  , muSampleSize :: Int
+  , muWeights    :: Weights
+  } deriving (Show, Eq, Generic)
+instance Binary ModelUpdate
+
+data Metrics = Metrics
+  { mAccuracy  :: Double
+  , mPrecision :: Double
+  , mRecall    :: Double
+  , mF1        :: Double
+  , mAUCROC    :: Double
+  , mRound     :: Int
+  } deriving (Show, Eq, Generic)
+instance Binary Metrics
+
+-- ---------------------------------------------------------------------------
+-- Logger
+-- ---------------------------------------------------------------------------
+
+data LogLevel = DEBUG | INFO | WARN | ERROR
+  deriving (Show, Eq, Ord, Generic)
+instance Binary LogLevel
+
+data LogEntry = LogEntry
+  { leSource :: String
+  , leLevel  :: LogLevel
+  , leMsg    :: String
+  } deriving (Show, Generic)
+instance Binary LogEntry
+
+data LogMsg = Log LogEntry
+  deriving (Show, Generic)
+instance Binary LogMsg
+
+-- ---------------------------------------------------------------------------
+-- Evaluator
+-- ---------------------------------------------------------------------------
+
+data EvaluatorMsg = EvaluateModel GlobalModel
+  deriving (Show, Generic)
+instance Binary EvaluatorMsg
+
+-- ---------------------------------------------------------------------------
+-- Trainer
+-- ---------------------------------------------------------------------------
+
+data TrainerMsg
+  = StartTraining GlobalModel
+  | TrainerStop
+  deriving (Show, Generic)
+instance Binary TrainerMsg
+
+-- ---------------------------------------------------------------------------
+-- Aggregator (provider mode)
+-- ---------------------------------------------------------------------------
+
+data AggregatorMsg = SubmitUpdate ModelUpdate
+  deriving (Show, Generic)
+instance Binary AggregatorMsg
+
+-- ---------------------------------------------------------------------------
+-- Coordinator (provider mode)
+-- ---------------------------------------------------------------------------
+
+data CoordinatorMsg
+  = RegisterTrainer String ActorId
+  | AggregationComplete GlobalModel
+  | EvaluationDone Metrics
+  | StartFederation
+  deriving (Show, Generic)
+instance Binary CoordinatorMsg
+
+-- ---------------------------------------------------------------------------
+-- PeerCoordinator (P2P mode)
+-- ---------------------------------------------------------------------------
+
+data PeerSyncPayload = PeerSyncPayload
+  { pspCounter      :: !GCounter
+  , pspParticipants :: !(ORSet String)
+  , pspUpdate       :: !(Maybe ModelUpdate)
+  } deriving (Show, Generic)
+instance Binary PeerSyncPayload
+
+data PeerCoordinatorMsg
+  = SetPeers ![ActorId]
+  | PeerStartRound
+  | PeerTrainingDone !ModelUpdate
+  | PeerSync !PeerSyncPayload
+  | PeerEvaluationDone !Metrics
+  deriving (Show, Generic)
+instance Binary PeerCoordinatorMsg