From 8ab0b946cbb24a25c4f2e65031584ea46c6a15d0 Mon Sep 17 00:00:00 2001
From: IC Rainbow <aenor.realm@gmail.com>
Date: Sun, 13 Aug 2023 14:24:44 +0300
Subject: [PATCH] Move custom bit of Process into ProcessEnv

Adds withProcessEnv (like withReaderT) - a more general `local` for MonadReader r (Process r).
---
 troupe/src/Troupe.hs         |  6 ++++--
 troupe/src/Troupe/Process.hs | 40 ++++++++++++++++--------------------
 2 files changed, 22 insertions(+), 24 deletions(-)

diff --git a/troupe/src/Troupe.hs b/troupe/src/Troupe.hs
index 0948e38..d80d232 100644
--- a/troupe/src/Troupe.hs
+++ b/troupe/src/Troupe.hs
@@ -68,6 +68,7 @@ module Troupe
 
     -- * @mtl@-style transformer support
     MonadProcess (..),
+    withProcessEnv
   )
 where
 
@@ -94,6 +95,7 @@ import Troupe.Process
   ( DemonitorOption (..),
     Match,
     MonadProcess (..),
+    withProcessEnv,
     NodeContext (..),
     Process,
     ProcessEnv (..),
@@ -132,9 +134,9 @@ import Troupe.Types (Down (..), MonitorRef, ProcessId)
 -- doesn't necessarily return when @p@ returns.
 runNode :: r -> Process r a -> IO ()
 runNode r process = do
-  nodeContext <- newNodeContext r
+  nodeContext <- newNodeContext
   processContext <- newProcessContext nodeContext
-  let processEnv = ProcessEnv nodeContext processContext
+  let processEnv = ProcessEnv nodeContext processContext r
 
   _ <- runProcess (spawn process) processEnv
 
diff --git a/troupe/src/Troupe/Process.hs b/troupe/src/Troupe/Process.hs
index d6c120c..f20cd89 100644
--- a/troupe/src/Troupe/Process.hs
+++ b/troupe/src/Troupe/Process.hs
@@ -21,6 +21,7 @@ module Troupe.Process
     MonadProcess (..),
     Process,
     runProcess,
+    withProcessEnv,
     self,
     ProcessOption (..),
     setProcessOption,
@@ -148,20 +149,18 @@ import Troupe.Types
     succProcessId,
   )
 
-data NodeContext r = NodeContext
+data NodeContext = NodeContext
   { nodeContextNextProcessId :: {-# UNPACK #-} !(TVar ProcessId),
     nodeContextNextMonitorRefId :: {-# UNPACK #-} !(TVar MonitorRefId),
-    nodeContextProcesses :: {-# UNPACK #-} !(Map ProcessId ProcessContext),
-    nodeContextR :: r
+    nodeContextProcesses :: {-# UNPACK #-} !(Map ProcessId ProcessContext)
   }
 
-newNodeContext :: r -> IO (NodeContext r)
-newNodeContext r =
+newNodeContext :: IO NodeContext
+newNodeContext =
   NodeContext
     <$> newTVarIO processId0
     <*> newTVarIO monitorRefId0
     <*> Map.newIO
-    <*> pure r
 
 data ProcessContext = ProcessContext
   { processContextId :: {-# UNPACK #-} !ProcessId,
@@ -173,7 +172,7 @@ data ProcessContext = ProcessContext
     processContextMonitorees :: {-# UNPACK #-} !(Set MonitorRef)
   }
 
-newProcessContextSTM :: CQueue Dynamic -> ReaderT (NodeContext r) STM ProcessContext
+newProcessContextSTM :: CQueue Dynamic -> ReaderT NodeContext STM ProcessContext
 newProcessContextSTM queue = do
   pid <- newPid
   lift $
@@ -189,15 +188,17 @@ newProcessContextSTM queue = do
       writeTVarR nodeContextNextProcessId $!! succProcessId curr
       pure curr
 
-newProcessContext :: NodeContext r -> IO ProcessContext
+newProcessContext :: NodeContext -> IO ProcessContext
 newProcessContext nodeContext = do
   queue <- newCQueue
   atomically $ runReaderT (newProcessContextSTM queue) nodeContext
 
 data ProcessEnv r = ProcessEnv
-  { processEnvNodeContext :: {-# UNPACK #-} !(NodeContext r),
-    processEnvProcessContext :: {-# UNPACK #-} !ProcessContext
+  { processEnvNodeContext :: {-# UNPACK #-} !NodeContext,
+    processEnvProcessContext :: {-# UNPACK #-} !ProcessContext,
+    processEnvLocalContext :: !r
   }
+  deriving stock (Functor)
 
 -- | @mtl@-style class to bring 'Process' support to transformers.
 class (Monad m) => MonadProcess r m | m -> r where
@@ -273,21 +274,16 @@ instance MonadProcess r (Process r) where
   {-# INLINE getProcessEnv #-}
 
 instance MonadReader r (Process r) where
-  ask = Process $ reader (nodeContextR . processEnvNodeContext)
-  reader f = Process $ reader (f . nodeContextR . processEnvNodeContext)
-  local f (Process a) = Process $ local mapR a
-    where
-      mapR env =
-        env
-          { processEnvNodeContext =
-              (processEnvNodeContext env)
-                { nodeContextR = f (nodeContextR (processEnvNodeContext env))
-                }
-          }
+  ask = Process $ reader processEnvLocalContext
+  reader f = Process $ reader (f . processEnvLocalContext)
+  local f (Process a) = Process $ local (fmap f) a
 
 runProcess :: Process r a -> ProcessEnv r -> IO a
 runProcess = runReaderT . unProcess
 
+withProcessEnv :: (r -> r') -> Process r' a -> Process r a
+withProcessEnv f (Process m) = Process $ withReaderT (fmap f) m
+
 -- | Get the 'ProcessId' of the running process.
 self :: (MonadProcess r m) => m ProcessId
 self = processContextId . processEnvProcessContext <$> getProcessEnv
@@ -360,7 +356,7 @@ deliverExit !pc !exc = do
       Nothing -> pure () -- Process exited without exception
       Just _ -> deliverException pc (toException exc')
 
-lookupProcess :: ProcessId -> ReaderT (NodeContext r) STM (Maybe ProcessContext)
+lookupProcess :: ProcessId -> ReaderT NodeContext STM (Maybe ProcessContext)
 lookupProcess !pid = reader nodeContextProcesses >>= lift . Map.lookup pid
 
 linkSTM :: ProcessId -> ReaderT (ProcessEnv r) STM ()