From 162d50d7ad3dc6c6933db57a7caa4dece3f655be Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Tue, 8 Apr 2025 08:10:08 -0700
Subject: [PATCH 01/16] MWB CPP pragma variable to handle custom GHC.

Change-Id: I68e30f82739d98aab253a3f927ed08b3233ae574
---
 cabal.project                             | 5 ++++-
 plugin/ghc-persistent-worker-plugin.cabal | 8 ++++++++
 plugin/src/GHCPersistentWorkerPlugin.hs   | 4 ++--
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/cabal.project b/cabal.project
index 0f150933..0f6a61a2 100644
--- a/cabal.project
+++ b/cabal.project
@@ -6,8 +6,11 @@ packages: ./plugin
           ./buck-worker
           ./buck-multiplex-worker
 
+-- For MWB-customized GHC, turn this on.
+-- flags: +mwb
+
 allow-newer:
-  proto-lens:base,
+  proto-lens:base
 
 if (impl(ghc >= 9.11))
   allow-newer:
diff --git a/plugin/ghc-persistent-worker-plugin.cabal b/plugin/ghc-persistent-worker-plugin.cabal
index b2f04681..78a1955c 100644
--- a/plugin/ghc-persistent-worker-plugin.cabal
+++ b/plugin/ghc-persistent-worker-plugin.cabal
@@ -16,6 +16,11 @@ extra-doc-files:    CHANGELOG.md
 common warnings
     ghc-options: -Wall
 
+flag mwb
+    description: use mwb-customized GHC
+    manual: True
+    default: False
+
 library
     import:           warnings
     hs-source-dirs:   src
@@ -43,6 +48,9 @@ library
     if (impl(ghc >= 9.8) && impl(ghc < 9.9))
         build-depends: base ^>=4.19
 
+    if (flag(mwb))
+        cpp-options: -DMWB
+
     build-depends:    binary,
                       bytestring,
                       containers,
diff --git a/plugin/src/GHCPersistentWorkerPlugin.hs b/plugin/src/GHCPersistentWorkerPlugin.hs
index 1e016f0c..7e0a3a6b 100644
--- a/plugin/src/GHCPersistentWorkerPlugin.hs
+++ b/plugin/src/GHCPersistentWorkerPlugin.hs
@@ -30,7 +30,7 @@ import System.Environment (setEnv)
 import System.IO (BufferMode (..), Handle, hFlush, hGetLine, hPutStrLn, hSetBuffering, stderr, stdin, stdout)
 import Text.Read (readMaybe)
 
-#if __GLASGOW_HASKELL__ >= 911
+#if __GLASGOW_HASKELL__ >= 911 || defined(MWB)
 import Control.Concurrent.MVar (newMVar)
 import GHC.Types.Unique.FM (emptyUFM)
 #endif
@@ -116,7 +116,7 @@ workerImplDefault =
       GHC.initGhcMonad Nothing
       -- explicitly initialize loader.
       loader <- liftIO Loader.uninitializedLoader
-#if __GLASGOW_HASKELL__ >= 911
+#if __GLASGOW_HASKELL__ >= 911 || defined(MWB)
       lookup_cache <- liftIO $ newMVar emptyUFM
       let interp = Interp InternalInterp loader lookup_cache
 #else

From 2953e3edb5b2474f19cc8e1315894b0c1147339c Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Tue, 8 Apr 2025 11:27:08 -0700
Subject: [PATCH 02/16] disable buck-multiplex-worker since the build is broken
 anyway.

Change-Id: Ib3b3646fd9c4db5180834a6a830cad6326e53ff2
---
 buck-multiplex-worker/Main.hs | 166 +++++++++++++++++-----------------
 1 file changed, 83 insertions(+), 83 deletions(-)

diff --git a/buck-multiplex-worker/Main.hs b/buck-multiplex-worker/Main.hs
index 14634fb6..8688482a 100644
--- a/buck-multiplex-worker/Main.hs
+++ b/buck-multiplex-worker/Main.hs
@@ -5,11 +5,10 @@ module Main where
 import BuckArgs (BuckArgs (..), CompileResult (..), parseBuckArgs, toGhcArgs, writeResult)
 import BuckWorker (
   ExecuteCommand (..),
-  ExecuteCommand_EnvironmentEntry (..),
+  ExecuteCommand'EnvironmentEntry (..),
   ExecuteEvent (..),
   ExecuteResponse (..),
   Worker (..),
-  workerServer,
   )
 import Control.Concurrent.MVar (MVar)
 import Control.Concurrent.STM (TVar, newTVarIO)
@@ -35,14 +34,6 @@ import Internal.Cache (Cache (..), ModuleArtifacts (..), emptyCache)
 import Internal.Log (newLog)
 import Internal.Session (Env (..), withGhc)
 import Message
-import Network.GRPC.HighLevel.Generated (
-  GRPCMethodType (..),
-  ServerRequest (..),
-  ServerResponse (..),
-  ServiceOptions (..),
-  StatusCode (..),
-  defaultServiceOptions,
-  )
 import Pool (Pool (..), dumpStatus, removeWorker)
 import Prelude hiding (log)
 import Server (assignLoop)
@@ -50,13 +41,14 @@ import System.Environment (lookupEnv)
 import System.IO (BufferMode (LineBuffering), hPutStrLn, hSetBuffering, stderr, stdout)
 import Worker (work)
 
-commandEnv :: Vector.Vector ExecuteCommand_EnvironmentEntry -> Map String String
-commandEnv =
-  Map.fromList .
-  fmap (\ (ExecuteCommand_EnvironmentEntry key value) -> (fromBs key, fromBs value)) .
-  Vector.toList
-  where
-    fromBs = Text.unpack . decodeUtf8Lenient
+commandEnv :: Vector.Vector ExecuteCommand'EnvironmentEntry -> Map String String
+commandEnv = undefined
+  -- THIS CODE IS NOT YET COMPATIBLE WITH grapesy yet.
+  -- Map.fromList .
+  -- fmap (\ (ExecuteCommand'EnvironmentEntry key value) -> (fromBs key, fromBs value)) .
+  --   Vector.toList
+  -- where
+  --   fromBs = Text.unpack . decodeUtf8Lenient
 
 abiHashIfSuccess :: Env -> BuckArgs -> Int -> IO (Maybe CompileResult)
 abiHashIfSuccess env args code
@@ -94,70 +86,78 @@ processRequest pool buckArgs env@Env {args} = do
       dumpStatus pool
       pure (result, unlines (responseConsoleStdOut ++ responseConsoleStdErr))
 
-executeHandler ::
-  MVar Cache ->
-  TVar Pool ->
-  ServerRequest 'Normal ExecuteCommand ExecuteResponse ->
-  IO (ServerResponse 'Normal ExecuteResponse)
-executeHandler cache pool (ServerNormalRequest _ ExecuteCommand {executeCommandArgv, executeCommandEnv}) = do
-  hPutStrLn stderr (unlines argv)
-  response <- either exceptionResponse successResponse =<< try run
-  pure (ServerNormalResponse response [] StatusOk "")
-  where
-    run = do
-      buckArgs <- either (throwIO . userError) pure (parseBuckArgs (commandEnv executeCommandEnv) argv)
-      args <- toGhcArgs buckArgs
-      log <- newLog False
-      result <- processRequest pool buckArgs Env {cache, args, log}
-      pure (buckArgs, result)
-
-    successResponse (buckArgs, (result, diagnostics)) = do
-      executeResponseExitCode <- writeResult buckArgs result
-      pure ExecuteResponse {
-        executeResponseExitCode,
-        executeResponseStderr = LazyText.pack diagnostics
-      }
-
-    exceptionResponse (SomeException e) =
-      pure ExecuteResponse {
-        executeResponseExitCode = 1,
-        executeResponseStderr = "Uncaught exception: " <> LazyText.pack (show e)
-      }
-
-    argv = Text.unpack . decodeUtf8Lenient <$> Vector.toList executeCommandArgv
-
-execHandler ::
-  ServerRequest 'ClientStreaming ExecuteEvent ExecuteResponse ->
-  IO (ServerResponse 'ClientStreaming ExecuteResponse)
-execHandler (ServerReaderRequest _metadata _recv) = do
-  hPutStrLn stderr "Received Exec"
-  error "not implemented"
-
-handlers :: MVar Cache -> TVar Pool -> Worker ServerRequest ServerResponse
-handlers cache srv =
-  Worker
-    { workerExecute = executeHandler cache srv,
-      workerExec = execHandler
-    }
-
 main :: IO ()
-main = do
-  hSetBuffering stdout LineBuffering
-  hSetBuffering stderr LineBuffering
-  socket <- lookupEnv "WORKER_SOCKET"
-  hPutStrLn stderr $ "using worker socket: " <> show socket
-  let
-    n = 1
-    thePool = Pool
-        { poolLimit = n,
-          poolNewWorkerId = 1,
-          poolNewJobId = 1,
-          poolStatus = mempty,
-          poolHandles = []
-        }
-
-  poolRef <- newTVarIO thePool
-  cache <- emptyCache False
-  workerServer (handlers cache poolRef) (maybe id setSocket socket defaultServiceOptions)
-  where
-    setSocket s options = options {serverHost = fromString ("unix://" <> s <> "\x00"), serverPort = 0}
+main = pure ()
+
+-- TODO: Revice the following code.
+--
+-- THE FOLLOWING CODE IS NOT COMPATIBLE WITH grapesy YET.
+--
+
+-- executeHandler ::
+--   MVar Cache ->
+--   TVar Pool ->
+--   ServerRequest 'Normal ExecuteCommand ExecuteResponse ->
+--   IO (ServerResponse 'Normal ExecuteResponse)
+-- executeHandler cache pool (ServerNormalRequest _ ExecuteCommand {executeCommandArgv, executeCommandEnv}) = do
+--   hPutStrLn stderr (unlines argv)
+--   response <- either exceptionResponse successResponse =<< try run
+--   pure (ServerNormalResponse response [] StatusOk "")
+--   where
+--     run = do
+--       buckArgs <- either (throwIO . userError) pure (parseBuckArgs (commandEnv executeCommandEnv) argv)
+--       args <- toGhcArgs buckArgs
+--       log <- newLog False
+--       result <- processRequest pool buckArgs Env {cache, args, log}
+--       pure (buckArgs, result)
+
+--     successResponse (buckArgs, (result, diagnostics)) = do
+--       executeResponseExitCode <- writeResult buckArgs result
+--       pure ExecuteResponse {
+--         executeResponseExitCode,
+--         executeResponseStderr = LazyText.pack diagnostics
+--       }
+
+--     exceptionResponse (SomeException e) =
+--       pure ExecuteResponse {
+--         executeResponseExitCode = 1,
+--         executeResponseStderr = "Uncaught exception: " <> LazyText.pack (show e)
+--       }
+
+--     argv = Text.unpack . decodeUtf8Lenient <$> Vector.toList executeCommandArgv
+
+-- execHandler ::
+--   ServerRequest 'ClientStreaming ExecuteEvent ExecuteResponse ->
+--   IO (ServerResponse 'ClientStreaming ExecuteResponse)
+-- execHandler (ServerReaderRequest _metadata _recv) = do
+--   hPutStrLn stderr "Received Exec"
+--   error "not implemented"
+
+-- handlers :: MVar Cache -> TVar Pool -> Worker ServerRequest ServerResponse
+-- handlers cache srv =
+--   Worker
+--     { workerExecute = executeHandler cache srv,
+--       workerExec = execHandler
+--     }
+
+-- main :: IO ()
+-- main = do
+--   hSetBuffering stdout LineBuffering
+--   hSetBuffering stderr LineBuffering
+--   socket <- lookupEnv "WORKER_SOCKET"
+--   hPutStrLn stderr $ "using worker socket: " <> show socket
+--   let
+--     n = 1
+--     thePool = Pool
+--         { poolLimit = n,
+--           poolNewWorkerId = 1,
+--           poolNewJobId = 1,
+--           poolStatus = mempty,
+--           poolHandles = []
+--         }
+
+--   poolRef <- newTVarIO thePool
+--   cache <- emptyCache False
+--   workerServer (handlers cache poolRef) (maybe id setSocket socket defaultServiceOptions)
+--   where
+--     setSocket s options = options {serverHost = fromString ("unix://" <> s <> "\x00"), serverPort = 0}

From 3e51bb2ee56b797aa92e58bd28d36cf162715d73 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Wed, 9 Apr 2025 01:02:47 -0700
Subject: [PATCH 03/16] shell.nix for minimal build shell.

Change-Id: I319d72bd7424cd79f01b5e82edfe2684da2dbfa5
---
 shell.nix | 9 +++++++++
 1 file changed, 9 insertions(+)
 create mode 100644 shell.nix

diff --git a/shell.nix b/shell.nix
new file mode 100644
index 00000000..67c430ec
--- /dev/null
+++ b/shell.nix
@@ -0,0 +1,9 @@
+let
+  pkgs = import (fetchTarball {
+    url = "https://github.com/nixos/nixpkgs/archive/e913ae340076bbb73d9f4d3d065c2bca7caafb16.tar.gz";
+    sha256 = "07qjibn85hc6p8c7lwg00fwpyqjlmaxncz9wa2l6qpy1hsk51k8f";
+  }) {};
+in pkgs.mkShell {
+  ghc_dir = "${pkgs.haskell.compiler.ghc910}";
+  packages = [pkgs.haskell.compiler.ghc910 pkgs.zlib pkgs.zlib.dev pkgs.snappy pkgs.protobuf pkgs.cabal-install];
+}

From 9613f88bea71cc8113b3cbf471e2775ca860557c Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Mon, 7 Apr 2025 14:50:27 -0700
Subject: [PATCH 04/16] Introduce WorkerMode

We invoke GHC session via oneshot mode or make mode in the worker

Change-Id: I20aaec6c737f8e3a56d90c6a4cfbe3aef827c082
---
 buck-worker/Main.hs | 53 ++++++++++++++++++++++++++++++---------------
 1 file changed, 35 insertions(+), 18 deletions(-)

diff --git a/buck-worker/Main.hs b/buck-worker/Main.hs
index 15f0d776..e5468407 100644
--- a/buck-worker/Main.hs
+++ b/buck-worker/Main.hs
@@ -2,6 +2,8 @@
 
 module Main where
 
+import BuckArgs (BuckArgs (abiOut), CompileResult (..), parseBuckArgs, toGhcArgs, writeResult)
+import BuckWorker (ExecuteCommand, ExecuteCommand'EnvironmentEntry, ExecuteEvent, ExecuteResponse, Worker (..), Instrument (..))
 import Control.Concurrent (MVar, modifyMVar_, newMVar)
 import Control.Concurrent.Async (async, cancel, wait)
 import Control.Concurrent.Chan (Chan, dupChan, readChan, newChan, writeChan)
@@ -43,8 +45,6 @@ import System.Exit (exitFailure)
 import System.FilePath (takeDirectory, (</>))
 import System.IO (BufferMode (LineBuffering), IOMode (..), hGetLine, hPutStrLn, hSetBuffering, stderr, stdout, hPutStr)
 
-import BuckArgs (BuckArgs (..), CompileResult (..), parseBuckArgs, toGhcArgs, writeResult)
-import BuckWorker (ExecuteCommand, ExecuteCommand'EnvironmentEntry, ExecuteEvent, ExecuteResponse, Worker (..), Instrument (..))
 import Proto.Worker_Fields qualified as Fields
 import Proto.Instrument_Fields qualified as Instr
 import qualified Proto.Instrument as Instr
@@ -54,6 +54,12 @@ data WorkerStatus =
     active :: Int
   }
 
+data WorkerMode =
+  WorkerMakeMode
+  |
+  WorkerOneshotMode
+  deriving stock (Eq, Show)
+
 commandEnv :: [Proto ExecuteCommand'EnvironmentEntry] -> Map String String
 commandEnv =
   Map.fromList .
@@ -61,12 +67,16 @@ commandEnv =
   where
     fromBs = Text.unpack . decodeUtf8Lenient
 
+-- | Compile a single module.
+-- Depending on @mode@ this will either use the old EPS-based oneshot-style compilation logic or the HPT-based
+-- make-style implementation.
 compileAndReadAbiHash ::
+  WorkerMode ->
   Chan (Proto Instr.Event) ->
   BuckArgs ->
   Target ->
   Ghc (Maybe CompileResult)
-compileAndReadAbiHash instrChan args target = do
+compileAndReadAbiHash _mode instrChan args target = do
 
   liftIO $ writeChan instrChan $
     defMessage &
@@ -128,10 +138,11 @@ debugRequestArgs = False
 execute ::
   MVar WorkerStatus ->
   MVar Cache ->
+  WorkerMode ->
   Chan (Proto Instr.Event) ->
   Proto ExecuteCommand ->
   IO (Proto ExecuteResponse)
-execute status cache instrChan req = do
+execute status cache mode instrChan req = do
   when debugRequestArgs do
     hPutStrLn stderr (unlines argv)
   msg <- either exceptionResponse successResponse =<< try run
@@ -157,7 +168,7 @@ execute status cache instrChan req = do
       bracket (startJob status) (finishJob status) \ _ -> do
         log <- newLog True
         let env = Env {log, cache, args}
-        result <- withGhc env (compileAndReadAbiHash instrChan buckArgs)
+        result <- withGhc env (compileAndReadAbiHash mode instrChan buckArgs)
         pure (env, buckArgs, result)
 
     successResponse (env, buckArgs, result) = do
@@ -193,12 +204,13 @@ exec _ =
 ghcServerMethods ::
   MVar WorkerStatus ->
   MVar Cache ->
+  WorkerMode ->
   Chan (Proto Instr.Event) ->
   Methods IO (ProtobufMethodsOf Worker)
-ghcServerMethods status cache instrChan =
+ghcServerMethods status cache mode instrChan =
   simpleMethods
     (mkClientStreaming exec)
-    (mkNonStreaming (execute status cache instrChan))
+    (mkNonStreaming (execute status cache mode instrChan))
 
 notifyMe ::
   Chan (Proto Instr.Event) ->
@@ -259,8 +271,8 @@ primarySocketDiscoveryIn :: FilePath -> PrimarySocketDiscoveryPath
 primarySocketDiscoveryIn dir = PrimarySocketDiscoveryPath (dir </> "primary")
 
 -- | Start a gRPC server that dispatches requests to GHC handlers.
-runLocalGhc :: ServerSocketPath -> Maybe InstrumentSocketPath -> IO ()
-runLocalGhc socket minstr = do
+runLocalGhc :: WorkerMode -> ServerSocketPath -> Maybe InstrumentSocketPath -> IO ()
+runLocalGhc mode socket minstr = do
   dbg ("Starting ghc server on " ++ socket.path)
 
   cache <- emptyCache True
@@ -271,12 +283,12 @@ runLocalGhc socket minstr = do
     dbg ("Instrumentation info available on " ++ instr.path)
     async $ runServerWithHandlers def (grpcServerConfig instr.path) $ fromMethods (instrumentMethods instrChan)
 
-  runServerWithHandlers def (grpcServerConfig socket.path) $ fromMethods (ghcServerMethods status cache instrChan)
+  runServerWithHandlers def (grpcServerConfig socket.path) $ fromMethods (ghcServerMethods status cache mode instrChan)
 
 -- | Start a gRPC server that runs GHC for client proxies, deleting the discovery file on shutdown.
-runCentralGhc :: PrimarySocketDiscoveryPath -> ServerSocketPath -> Maybe InstrumentSocketPath -> IO ()
-runCentralGhc discovery socket instr =
-  finally (runLocalGhc socket instr) do
+runCentralGhc :: WorkerMode -> PrimarySocketDiscoveryPath -> ServerSocketPath -> Maybe InstrumentSocketPath -> IO ()
+runCentralGhc mode discovery socket instr =
+  finally (runLocalGhc mode socket instr) do
     dbg ("Shutting down ghc server on " ++ socket.path)
     removeFile discovery.path
 
@@ -350,7 +362,8 @@ runOrProxyCentralGhc socket = do
       Right !primary | not (null primary) -> do
         pure (Left (PrimarySocketPath primary))
       _ -> do
-        thread <- async (runCentralGhc primaryFile socket instr)
+        let mode = WorkerOneshotMode
+        thread <- async (runCentralGhc mode primaryFile socket instr)
         hPutStr handle socket.path
         pure (Right thread)
   case result of
@@ -366,24 +379,28 @@ runOrProxyCentralGhc socket = do
 data CliOptions =
   CliOptions {
     -- | Should only a single central GHC server be run, with all other worker processes proxying it?
-    single :: Bool
+    single :: Bool,
+
+    -- | The worker implementation: Make mode or oneshot mode.
+    mode :: WorkerMode
   }
   deriving stock (Eq, Show)
 
 defaultCliOptions :: CliOptions
-defaultCliOptions = CliOptions {single = False}
+defaultCliOptions = CliOptions {single = False, mode = WorkerOneshotMode}
 
 parseOptions :: [String] -> IO CliOptions
 parseOptions =
   flip foldM defaultCliOptions \ z -> \case
     "--single" -> pure z {single = True}
+    "--make" -> pure z {mode = WorkerMakeMode}
     arg -> throwIO (userError ("Invalid worker CLI arg: " ++ arg))
 
 runWorker :: ServerSocketPath -> CliOptions -> IO ()
-runWorker socket CliOptions {single} =
+runWorker socket CliOptions {single, mode} =
   if single
   then runOrProxyCentralGhc socket
-  else runLocalGhc socket Nothing
+  else runLocalGhc mode socket Nothing
 
 main :: IO ()
 main = do

From ad01fc07ab80db4bfd818f8f1316b95fe68833b6 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Thu, 10 Apr 2025 15:57:04 -0700
Subject: [PATCH 05/16] Remove outdated ghc-persistent-worker-client

Change-Id: I0d0b52e8e61285c3e74279333e6b309f88d7feb1
---
 cabal.project                             |   1 -
 client/CHANGELOG.md                       |   0
 client/LICENSE                            |   0
 client/app/Client.hs                      | 100 ----------------------
 client/ghc-persistent-worker-client.cabal |  27 ------
 5 files changed, 128 deletions(-)
 delete mode 100644 client/CHANGELOG.md
 delete mode 100644 client/LICENSE
 delete mode 100644 client/app/Client.hs
 delete mode 100644 client/ghc-persistent-worker-client.cabal

diff --git a/cabal.project b/cabal.project
index 0f6a61a2..dff807ad 100644
--- a/cabal.project
+++ b/cabal.project
@@ -1,7 +1,6 @@
 packages: ./plugin
           ./comm
           ./server
-          ./client
           ./instrument
           ./buck-worker
           ./buck-multiplex-worker
diff --git a/client/CHANGELOG.md b/client/CHANGELOG.md
deleted file mode 100644
index e69de29b..00000000
diff --git a/client/LICENSE b/client/LICENSE
deleted file mode 100644
index e69de29b..00000000
diff --git a/client/app/Client.hs b/client/app/Client.hs
deleted file mode 100644
index d84155e6..00000000
--- a/client/app/Client.hs
+++ /dev/null
@@ -1,100 +0,0 @@
-{-# LANGUAGE OverloadedStrings #-}
-module Main (main) where
-
-import qualified Control.Exception as E
-import qualified Data.ByteString.Char8 as C
-import qualified Data.ByteString.Lazy as L
-import Data.Binary (encode)
-import Data.Foldable (foldMap)
-import Data.Int (Int32)
-import Data.List (isPrefixOf, lookup, partition, stripPrefix)
-import Data.Monoid (First (..))
-import Message
-  ( TargetId (..),
-    Msg (..),
-    Request (..),
-    Response (..),
-    recvMsg,
-    sendMsg,
-    unwrapMsg,
-    wrapMsg
-  )
-import Network.Socket
-  ( Family (AF_UNIX),
-    SockAddr (SockAddrUnix),
-    Socket,
-    SocketType (Stream),
-    close,
-    connect,
-    socket,
-    withSocketsDo,
-  )
-import Network.Socket.ByteString (recv, sendAll)
-import System.Environment (getArgs, getEnv, getEnvironment)
-import System.Exit (exitFailure)
-import System.IO (hFlush, hPutStrLn, stderr, stdout)
-
-data WorkerConfig = WorkerConfig
-  { workerConfigSocket :: String,
-    workerConfigTargetId :: Maybe TargetId,
-    workerConfigClose :: Bool
-  }
-  deriving (Show)
-
-splitArgs :: [String] -> ([String], [String])
-splitArgs = partition ("--worker-" `isPrefixOf`)
-
-getWorkerConfig :: [String] -> Maybe WorkerConfig
-getWorkerConfig args = do
-  socket <- getFirst $ foldMap (First . stripPrefix "--worker-socket=") args
-  let mid = getFirst $ foldMap (First . stripPrefix "--worker-target-id=") args
-      willClose = any ("--worker-close" `isPrefixOf`) args
-  pure WorkerConfig
-    { workerConfigSocket = socket,
-      workerConfigTargetId = TargetId <$> mid,
-      workerConfigClose = willClose
-    }
-
-main :: IO ()
-main = do
-  args <- getArgs
-  let (workerArgs, ghcArgs) = splitArgs args
-      mConf = getWorkerConfig workerArgs
-  hPutStrLn stderr (show mConf)
-  hPutStrLn stderr (show args)
-  hFlush stderr
-  case mConf of
-    Nothing -> do
-      hPutStrLn stderr "ghc-persistent-worker-client: Please pass --worker-socket=(socket file path)."
-      exitFailure
-    Just conf -> do
-      let sockPath = workerConfigSocket conf
-          mid = workerConfigTargetId conf
-          willClose = workerConfigClose conf
-      env <- getEnvironment
-      process sockPath mid willClose env ghcArgs
-
-process :: FilePath -> Maybe TargetId -> Bool -> [(String, String)] -> [String] -> IO ()
-process socketPath mid willClose env args = runClient socketPath $ \s -> do
-  let req = Request
-        { requestWorkerTargetId = mid,
-          requestWorkerClose = willClose,
-          requestEnv = env,
-          requestArgs = args
-        }
-  let msg = wrapMsg req
-  sendMsg s msg
-  --
-  msg' <- recvMsg s
-  let Response res ss_out ss_err = unwrapMsg msg'
-  mapM_ (hPutStrLn stdout) ss_out
-  hFlush stdout
-  mapM_ (hPutStrLn stderr) ss_err
-  hFlush stderr
-
-runClient :: FilePath -> (Socket -> IO a) -> IO a
-runClient fp client = withSocketsDo $ E.bracket (open fp) close client
-  where
-    open fp = E.bracketOnError (socket AF_UNIX Stream 0) close $ \sock -> do
-        connect sock (SockAddrUnix fp)
-        return sock
diff --git a/client/ghc-persistent-worker-client.cabal b/client/ghc-persistent-worker-client.cabal
deleted file mode 100644
index d77fa20c..00000000
--- a/client/ghc-persistent-worker-client.cabal
+++ /dev/null
@@ -1,27 +0,0 @@
-cabal-version:      3.4
-name:               ghc-persistent-worker-client
-version:            0.1.0.0
--- synopsis:
--- description:
-license:            MIT
-license-file:       LICENSE
-author:             Ian-Woo Kim
-maintainer:         ianwookim@gmail.com
--- copyright:
-category:           Development
-build-type:         Simple
-extra-doc-files:    CHANGELOG.md
--- extra-source-files:
-
-common warnings
-    ghc-options: -Wall
-
-executable Client
-    main-is:          Client.hs
-    hs-source-dirs:   app
-    build-depends:    base,
-                      binary,
-                      bytestring,
-                      ghc-persistent-worker-comm,
-                      network
-    default-language: GHC2021

From 0d16f39d451efdaa9ed18fe2e72d3c3405bd5ea5 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Thu, 10 Apr 2025 16:25:22 -0700
Subject: [PATCH 06/16] Remove comm and server.

Merged them into buck-multiplex-worker

Change-Id: I8e8f92b910ddd4e189bbfa766a7f8a5840814a5d
---
 buck-multiplex-worker/BUCK                    |  2 +-
 .../src => buck-multiplex-worker}/Message.hs  |  0
 {server/lib => buck-multiplex-worker}/Pool.hs |  0
 .../lib => buck-multiplex-worker}/Server.hs   |  0
 {comm/src => buck-multiplex-worker}/Util.hs   |  0
 .../lib => buck-multiplex-worker}/Worker.hs   |  0
 .../buck-multiplex-worker.cabal               | 19 ++++++-
 cabal.project                                 |  2 -
 comm/CHANGELOG.md                             |  0
 comm/LICENSE                                  |  0
 comm/ghc-persistent-worker-comm.cabal         | 35 ------------
 comm/src/BUCK                                 |  1 -
 server/CHANGELOG.md                           |  0
 server/LICENSE                                |  0
 server/app/Main.hs                            | 54 ------------------
 server/ghc-persistent-worker-server.cabal     | 56 -------------------
 server/lib/BUCK                               |  1 -
 17 files changed, 18 insertions(+), 152 deletions(-)
 rename {comm/src => buck-multiplex-worker}/Message.hs (100%)
 rename {server/lib => buck-multiplex-worker}/Pool.hs (100%)
 rename {server/lib => buck-multiplex-worker}/Server.hs (100%)
 rename {comm/src => buck-multiplex-worker}/Util.hs (100%)
 rename {server/lib => buck-multiplex-worker}/Worker.hs (100%)
 delete mode 100644 comm/CHANGELOG.md
 delete mode 100644 comm/LICENSE
 delete mode 100644 comm/ghc-persistent-worker-comm.cabal
 delete mode 100644 comm/src/BUCK
 delete mode 100644 server/CHANGELOG.md
 delete mode 100644 server/LICENSE
 delete mode 100644 server/app/Main.hs
 delete mode 100644 server/ghc-persistent-worker-server.cabal
 delete mode 100644 server/lib/BUCK

diff --git a/buck-multiplex-worker/BUCK b/buck-multiplex-worker/BUCK
index 0d9534ef..823fb11e 100644
--- a/buck-multiplex-worker/BUCK
+++ b/buck-multiplex-worker/BUCK
@@ -1 +1 @@
-export_file("Main.hs", visibility = ["PUBLIC"])
+[export_file(f + ".hs", visibility = ["PUBLIC"]) for f in ["Main", "Message", "Pool", "Server", "Util", "Worker"]]
diff --git a/comm/src/Message.hs b/buck-multiplex-worker/Message.hs
similarity index 100%
rename from comm/src/Message.hs
rename to buck-multiplex-worker/Message.hs
diff --git a/server/lib/Pool.hs b/buck-multiplex-worker/Pool.hs
similarity index 100%
rename from server/lib/Pool.hs
rename to buck-multiplex-worker/Pool.hs
diff --git a/server/lib/Server.hs b/buck-multiplex-worker/Server.hs
similarity index 100%
rename from server/lib/Server.hs
rename to buck-multiplex-worker/Server.hs
diff --git a/comm/src/Util.hs b/buck-multiplex-worker/Util.hs
similarity index 100%
rename from comm/src/Util.hs
rename to buck-multiplex-worker/Util.hs
diff --git a/server/lib/Worker.hs b/buck-multiplex-worker/Worker.hs
similarity index 100%
rename from server/lib/Worker.hs
rename to buck-multiplex-worker/Worker.hs
diff --git a/buck-multiplex-worker/buck-multiplex-worker.cabal b/buck-multiplex-worker/buck-multiplex-worker.cabal
index 91ba09f2..7df5a8c9 100644
--- a/buck-multiplex-worker/buck-multiplex-worker.cabal
+++ b/buck-multiplex-worker/buck-multiplex-worker.cabal
@@ -8,18 +8,33 @@ build-type: Simple
 
 executable multiplex-worker
   main-is: Main.hs
+  other-modules:
+    Message
+    Pool
+    Server
+    Util
+    Worker
   hs-source-dirs: .
   ghc-options: -Wall -Widentities -Wincomplete-uni-patterns -Wmissing-deriving-strategies -Wredundant-constraints -Wunused-type-patterns -Wunused-packages -threaded -rtsopts -with-rtsopts=-N
   build-depends:
     base ==4.*,
+    binary,
+    bytestring,
     containers,
+    deepseq,
+    directory,
+    extra,
+    filepath,
     ghc,
-    ghc-persistent-worker-comm,
     ghc-persistent-worker-plugin,
-    ghc-persistent-worker-server,
+    network,
+    optparse-applicative,
+    process,
     stm,
     text,
+    time,
     transformers,
+    unix,
     vector,
     buck-worker
   default-language: GHC2021
diff --git a/cabal.project b/cabal.project
index dff807ad..404d2cef 100644
--- a/cabal.project
+++ b/cabal.project
@@ -1,6 +1,4 @@
 packages: ./plugin
-          ./comm
-          ./server
           ./instrument
           ./buck-worker
           ./buck-multiplex-worker
diff --git a/comm/CHANGELOG.md b/comm/CHANGELOG.md
deleted file mode 100644
index e69de29b..00000000
diff --git a/comm/LICENSE b/comm/LICENSE
deleted file mode 100644
index e69de29b..00000000
diff --git a/comm/ghc-persistent-worker-comm.cabal b/comm/ghc-persistent-worker-comm.cabal
deleted file mode 100644
index bbaca00a..00000000
--- a/comm/ghc-persistent-worker-comm.cabal
+++ /dev/null
@@ -1,35 +0,0 @@
-cabal-version:      3.4
-name:               ghc-persistent-worker-comm
-version:            0.1.0.0
--- synopsis:
--- description:
-license:            MIT
-license-file:       LICENSE
-author:             Ian-Woo Kim
-maintainer:         ianwookim@gmail.com
--- copyright:
-category:           Development
-build-type:         Simple
-extra-doc-files:    CHANGELOG.md
--- extra-source-files:
-
-common warnings
-    ghc-options: -Wall
-
-library
-    import:           warnings
-    exposed-modules:  Message
-                      Util
-
-    build-depends:    base >=4.19,
-                      binary,
-                      bytestring,
-                      containers,
-                      deepseq,
-                      filepath,
-                      network,
-                      time,
-                      transformers,
-                      unix
-    hs-source-dirs:   src
-    default-language: GHC2021
diff --git a/comm/src/BUCK b/comm/src/BUCK
deleted file mode 100644
index 2473bae0..00000000
--- a/comm/src/BUCK
+++ /dev/null
@@ -1 +0,0 @@
-[export_file(f, visibility = ["PUBLIC"]) for f in ["Message.hs"]]
diff --git a/server/CHANGELOG.md b/server/CHANGELOG.md
deleted file mode 100644
index e69de29b..00000000
diff --git a/server/LICENSE b/server/LICENSE
deleted file mode 100644
index e69de29b..00000000
diff --git a/server/app/Main.hs b/server/app/Main.hs
deleted file mode 100644
index ee22f916..00000000
--- a/server/app/Main.hs
+++ /dev/null
@@ -1,54 +0,0 @@
-module Main where
-
-import Control.Concurrent.STM (newTVarIO)
-import Control.Monad (replicateM_)
-import qualified Data.IntMap as IM
-import Options.Applicative (Parser, (<**>))
-import qualified Options.Applicative as OA
-import Pool (Pool (..))
-import Server (runServer, serve, spawnWorker)
-
--- cli args
-data Option = Option
-  { optionNumWorkers :: Int,
-    optionGHC :: FilePath,
-    optionSocket :: FilePath,
-    optionPkgDbs :: [FilePath]
-  }
-
-p_option :: Parser Option
-p_option =
-  Option
-    <$> OA.option OA.auto
-        ( OA.long "num"
-            <> OA.short 'n'
-            <> OA.help "number of workers"
-        )
-    <*> OA.strOption
-        ( OA.long "ghc"
-            <> OA.help "GHC path"
-        )
-    <*> OA.strOption
-        ( OA.long "socket-file"
-            <> OA.help "UNIX socket file accepting compilation requests"
-        )
-    <*> OA.many (OA.strOption (OA.long "package-db" <> OA.help "Package DB Path"))
-
-main :: IO ()
-main = do
-  opts <- OA.execParser (OA.info (p_option <**> OA.helper) OA.fullDesc)
-  let n = optionNumWorkers opts
-      ghcPath = optionGHC opts
-      socketPath = optionSocket opts
-      dbPaths = optionPkgDbs opts
-  let thePool = Pool
-        { poolLimit = n,
-          poolNewWorkerId = 1,
-          poolNewJobId = 1,
-          poolStatus = IM.empty,
-          poolHandles = []
-        }
-
-  poolRef <- newTVarIO thePool
-  replicateM_ n $ spawnWorker False ghcPath dbPaths poolRef
-  runServer socketPath (serve ghcPath dbPaths poolRef)
diff --git a/server/ghc-persistent-worker-server.cabal b/server/ghc-persistent-worker-server.cabal
deleted file mode 100644
index 1c2d5607..00000000
--- a/server/ghc-persistent-worker-server.cabal
+++ /dev/null
@@ -1,56 +0,0 @@
-cabal-version:      3.4
-name:               ghc-persistent-worker-server
-version:            0.1.0.0
--- synopsis:
--- description:
-license:            MIT
-license-file:       LICENSE
-author:             Ian-Woo Kim
-maintainer:         ianwookim@gmail.com
--- copyright:
-category:           Development
-build-type:         Simple
-extra-doc-files:    CHANGELOG.md
-
-library
-    ghc-options:      -Wall -Werror
-    hs-source-dirs:   lib
-    exposed-modules:  Pool
-                      Server
-                      Worker
-    build-depends:    base >=4.19,
-                      binary,
-                      bytestring,
-                      containers,
-                      directory,
-                      extra,
-                      filepath,
-                      ghc-persistent-worker-comm,
-                      network,
-                      optparse-applicative,
-                      process,
-                      stm,
-                      transformers,
-                      unix
-    default-language: GHC2021
-
-executable Server
-    main-is:          Main.hs
-    ghc-options:      -Wall -Werror
-    hs-source-dirs:   app
-    build-depends:    base >=4.19,
-                      binary,
-                      bytestring,
-                      containers,
-                      directory,
-                      extra,
-                      filepath,
-                      ghc-persistent-worker-comm,
-                      ghc-persistent-worker-server,
-                      network,
-                      optparse-applicative,
-                      process,
-                      stm,
-                      transformers,
-                      unix
-    default-language: GHC2021
diff --git a/server/lib/BUCK b/server/lib/BUCK
deleted file mode 100644
index 021491be..00000000
--- a/server/lib/BUCK
+++ /dev/null
@@ -1 +0,0 @@
-[export_file(f + ".hs", visibility = ["PUBLIC"]) for f in ["Server", "Pool", "Worker"]]

From d4b5710182150a33aa23971f7db13937e5d2a760 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Mon, 7 Apr 2025 15:09:33 -0700
Subject: [PATCH 07/16] Store hug and moduleGraph in Cache. update them
 properly

Change-Id: Iee49bc4bfe2d000dde9ecee4cab958e124de3195
---
 buck-worker/Main.hs          |  17 +++-
 plugin/src/Internal/Cache.hs | 148 ++++++++++++++++++++++++++---------
 2 files changed, 124 insertions(+), 41 deletions(-)

diff --git a/buck-worker/Main.hs b/buck-worker/Main.hs
index e5468407..665cb9ec 100644
--- a/buck-worker/Main.hs
+++ b/buck-worker/Main.hs
@@ -28,7 +28,7 @@ import GHC (Ghc, getSession)
 import GHC.IO.Handle.FD (withFileBlocking)
 import GHC.Stats (getRTSStats, RTSStats (..), GCDetails (..))
 import Internal.AbiHash (AbiHash (..), showAbiHash)
-import Internal.Cache (Cache (..), ModuleArtifacts (..), Target (..), emptyCache)
+import Internal.Cache (Cache (..), CacheFeatures (..), ModuleArtifacts (..), Target (..), emptyCache, emptyCacheWith)
 import Internal.Compile (compile)
 import Internal.Log (dbg, logFlush, newLog)
 import Internal.Session (Env (..), withGhc)
@@ -275,7 +275,18 @@ runLocalGhc :: WorkerMode -> ServerSocketPath -> Maybe InstrumentSocketPath -> I
 runLocalGhc mode socket minstr = do
   dbg ("Starting ghc server on " ++ socket.path)
 
-  cache <- emptyCache True
+  cache <-
+    case mode of
+      WorkerMakeMode ->
+        emptyCacheWith CacheFeatures {
+          hpt = True,
+          loader = False,
+          enable = True,
+          names = False,
+          finder = True,
+          eps = False
+        }
+      WorkerOneshotMode -> emptyCache True
   status <- newMVar WorkerStatus {active = 0}
   instrChan <- newChan
 
@@ -412,5 +423,5 @@ main = do
   try (runWorker socket options) >>= \case
     Right () ->
       dbg "Worker terminated without cancellation."
-    Left (err :: SomeException) -> do
+    Left (err :: SomeException) ->
       dbg ("Worker terminated with exception: " ++ displayException err)
diff --git a/plugin/src/Internal/Cache.hs b/plugin/src/Internal/Cache.hs
index 0f87a9d2..87fe2bee 100644
--- a/plugin/src/Internal/Cache.hs
+++ b/plugin/src/Internal/Cache.hs
@@ -3,7 +3,7 @@
 module Internal.Cache where
 
 import Control.Concurrent.MVar (MVar, modifyMVar, modifyMVar_, newMVar, readMVar)
-import Control.Monad (join, unless)
+import Control.Monad (join, unless, when)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Data.Bifunctor (first)
 import Data.Coerce (coerce)
@@ -28,7 +28,7 @@ import GHC.Types.Name.Cache (NameCache (..), OrigNameCache)
 import GHC.Types.Unique.DFM (plusUDFM)
 import GHC.Types.Unique.FM (UniqFM, minusUFM, nonDetEltsUFM, sizeUFM)
 import GHC.Types.Unique.Supply (initUniqSupply)
-import GHC.Unit.Env (UnitEnv (..))
+import GHC.Unit.Env (HomeUnitEnv (..), HomeUnitGraph, UnitEnv (..), unitEnv_union)
 import GHC.Unit.External (ExternalUnitCache (..), initExternalUnitCache)
 import GHC.Unit.Module.Env (emptyModuleEnv, moduleEnvKeys, plusModuleEnv)
 import qualified GHC.Utils.Outputable as Outputable
@@ -36,14 +36,15 @@ import GHC.Utils.Outputable (SDoc, comma, doublePrec, fsep, hang, nest, punctuat
 import Internal.Log (Log, logd)
 import System.Environment (lookupEnv)
 
-#if __GLASGOW_HASKELL__ >= 911
+#if __GLASGOW_HASKELL__ >= 911 || defined(MWB)
 
 import Data.IORef (IORef, newIORef, readIORef)
 import qualified Data.Map.Lazy as LazyMap
 import GHC.Fingerprint (Fingerprint, getFileHash)
 import GHC.IORef (atomicModifyIORef')
 import GHC.Unit (InstalledModule, extendInstalledModuleEnv, lookupInstalledModuleEnv)
-import GHC.Unit.Finder (FinderCache (..), InstalledFindResult (..))
+import GHC.Unit.Finder (InstalledFindResult (..))
+import GHC.Unit.Finder.Types (FinderCache (..))
 import GHC.Unit.Module.Env (InstalledModuleEnv, emptyInstalledModuleEnv)
 import GHC.Utils.Panic (panic)
 
@@ -53,6 +54,14 @@ import GHC.Unit.Finder (FinderCache, initFinderCache)
 
 #endif
 
+import GHC.Unit.Module.Graph (ModuleGraph, unionMG)
+
+#if defined(MWB)
+
+import GHC.Unit.Module.Graph (ModuleGraphNode (..), mgModSummaries', mkModuleGraph, mkNodeKey)
+
+#endif
+
 data ModuleArtifacts =
   ModuleArtifacts {
     iface :: ModIface,
@@ -175,7 +184,7 @@ data BinPath =
   }
   deriving stock (Eq, Show)
 
-#if __GLASGOW_HASKELL__ >= 911
+#if __GLASGOW_HASKELL__ >= 911 || defined(MWB)
 
 data FinderState =
   FinderState {
@@ -210,35 +219,37 @@ data CacheFeatures =
     enable :: Bool,
     loader :: Bool,
     names :: Bool,
-    finder :: Bool
+    finder :: Bool,
+    eps :: Bool,
+    hpt :: Bool
   }
   deriving stock (Eq, Show)
 
 newCacheFeatures :: CacheFeatures
-newCacheFeatures = CacheFeatures {enable = True, loader = True, names = True, finder = True}
+newCacheFeatures = CacheFeatures {enable = True, loader = True, names = True, finder = True, eps = True, hpt = False}
 
 -- TODO the name cache could in principle be shared directly – try it out
 data Cache =
   Cache {
     features :: CacheFeatures,
-    initialized :: Bool,
     interp :: Maybe InterpCache,
     names :: OrigNameCache,
     stats :: Map Target CacheStats,
     path :: BinPath,
     finder :: FinderState,
     eps :: ExternalUnitCache,
+    hug :: Maybe HomeUnitGraph,
+    moduleGraph :: Maybe ModuleGraph,
     baseSession :: Maybe HscEnv
   }
 
-emptyCache :: Bool -> IO (MVar Cache)
-emptyCache enable = do
+emptyCacheWith :: CacheFeatures -> IO (MVar Cache)
+emptyCacheWith features = do
   initialPath <- lookupEnv "PATH"
   finder <- emptyFinderState
   eps <- initExternalUnitCache
   newMVar Cache {
-    features = newCacheFeatures {enable},
-    initialized = False,
+    features,
     interp = Nothing,
     names = emptyModuleEnv,
     stats = mempty,
@@ -248,14 +259,14 @@ emptyCache enable = do
     },
     finder,
     eps,
+    hug = Nothing,
+    moduleGraph = Nothing,
     baseSession = Nothing
   }
 
-initialize :: Cache -> IO Cache
-initialize cache = do
-  unless cache.initialized do
-    initUniqSupply 0 1
-  pure cache {initialized = True}
+emptyCache :: Bool -> IO (MVar Cache)
+emptyCache enable = do
+  emptyCacheWith newCacheFeatures {enable}
 
 basicLinkerStats :: LinkerEnv -> LinkerEnv -> LinkerStats
 basicLinkerStats base update =
@@ -521,7 +532,7 @@ report logVar workerId target cache = do
 
     workerDesc wid = text (" (" ++ wid ++ ")")
 
-#if __GLASGOW_HASKELL__ >= 911
+#if __GLASGOW_HASKELL__ >= 911 || defined(MWB)
 
 -- | This replacement of the Finder implementation has the sole purpose of recording some cache stats, for now.
 -- While its mutable state is allocated separately and shared across sessions, this doesn't really make a difference at
@@ -580,7 +591,7 @@ newFinderCache _ Cache {finder = FinderState {cache}} _ = pure cache
 
 withHscState :: HscEnv -> (MVar OrigNameCache -> MVar (Maybe LoaderState) -> MVar SymbolMap -> IO a) -> IO (Maybe a)
 withHscState HscEnv {hsc_interp, hsc_NC = NameCache {nsNames}} use =
-#if __GLASGOW_HASKELL__ >= 911
+#if __GLASGOW_HASKELL__ >= 911 || defined(MWB)
   for hsc_interp \ Interp {interpLoader = Loader {loader_state}, interpLookupSymbolCache} ->
     liftIO $ use nsNames loader_state interpLookupSymbolCache
 #else
@@ -590,33 +601,85 @@ withHscState HscEnv {hsc_interp, hsc_NC = NameCache {nsNames}} use =
     use nsNames loader_state symbolCacheVar
 #endif
 
+mergeHugs ::
+  HomeUnitEnv ->
+  HomeUnitEnv ->
+  HomeUnitEnv
+mergeHugs old new =
+  new {homeUnitEnv_hpt = plusUDFM old.homeUnitEnv_hpt new.homeUnitEnv_hpt}
+
 setTarget :: MVar Cache -> Cache -> HscEnv -> Target -> IO HscEnv
 setTarget cacheVar cache hsc_env target = do
-  hsc_FC <- newFinderCache cacheVar cache target
-  pure hsc_env {hsc_FC, hsc_unit_env = hsc_env.hsc_unit_env {ue_eps = cache.eps}}
+  hsc_env1 <-
+    if cache.features.finder
+    then do
+      hsc_FC <- newFinderCache cacheVar cache target
+      pure hsc_env {hsc_FC}
+    else pure hsc_env
+  let hsc_env2 =
+        if cache.features.eps
+        then hsc_env1 {hsc_unit_env = hsc_env1.hsc_unit_env {ue_eps = cache.eps}}
+        else hsc_env1
+      hsc_env3 =
+        maybe hsc_env2 (\ hug -> hsc_env2 {hsc_unit_env = hsc_env2.hsc_unit_env {ue_home_unit_graph = unitEnv_union mergeHugs hug hsc_env.hsc_unit_env.ue_home_unit_graph}}) cache.hug
+      hsc_env4 =
+        maybe id restoreModuleGraph cache.moduleGraph hsc_env3
+  pure hsc_env4
+  where
+    restoreModuleGraph mg e = e {hsc_mod_graph = mg}
+
+updateModuleGraph :: MVar Cache -> ModuleGraph -> IO ()
+updateModuleGraph cacheVar new =
+  modifyMVar_ cacheVar \ cache -> do
+#if defined(MWB)
+    pure cache {moduleGraph = Just (maybe new merge cache.moduleGraph)}
+  where
+    merge old =
+      mkModuleGraph (Map.elems (Map.unionWith mergeNodes oldMap newMap))
+      where
+        mergeNodes = \cases
+          (ModuleNode oldDeps _) (ModuleNode newDeps summ) -> ModuleNode (mergeDeps oldDeps newDeps) summ
+          _ newNode -> newNode
+
+        mergeDeps oldDeps newDeps = Set.toList (Set.fromList oldDeps <> Set.fromList newDeps)
+
+        oldMap = Map.fromList $ [(mkNodeKey n, n) | n <- mgModSummaries' old]
+
+        newMap = Map.fromList $ [(mkNodeKey n, n) | n <- mgModSummaries' new]
+#else
+    pure cache {moduleGraph = Just (maybe id unionMG cache.moduleGraph new)}
+#endif
 
 prepareCache :: MVar Cache -> Target -> HscEnv -> Cache -> IO (Cache, (HscEnv, Bool))
 prepareCache cacheVar target hsc_env0 cache0 = do
-  cache1 <- initialize cache0
   result <-
-    if cache1.features.enable
+    if cache0.features.enable
     then do
       hsc_env1 <- setTarget cacheVar cache0 hsc_env0 target
-      withHscState hsc_env1 \ nsNames loaderStateVar symbolCacheVar -> do
-        cache2 <- modifyMVar loaderStateVar \ initialLoaderState ->
-          modifyMVar symbolCacheVar \ initialSymbolCache ->
-            (first coerce) <$>
-            modifyMVar nsNames \ names ->
-              restoreCache target initialLoaderState (SymbolCache initialSymbolCache) names cache1
-        pure (hsc_env1, cache2)
+      if cache0.features.loader
+      then do
+        withHscState hsc_env1 \ nsNames loaderStateVar symbolCacheVar -> do
+          cache1 <- modifyMVar loaderStateVar \ initialLoaderState ->
+            modifyMVar symbolCacheVar \ initialSymbolCache ->
+              (first coerce) <$>
+              modifyMVar nsNames \ names ->
+                restoreCache target initialLoaderState (SymbolCache initialSymbolCache) names cache0
+          pure (hsc_env1, cache1)
+      else pure (Just (hsc_env1, cache0))
     else pure Nothing
-  let (hsc_env1, cache2) = fromMaybe (hsc_env0, cache1 {features = cache1.features {enable = False}}) result
-  pure (cache2, (hsc_env1, cache2.features.enable))
+  let (hsc_env1, cache1) = fromMaybe (hsc_env0, cache0 {features = cache0.features {loader = False}}) result
+  pure (cache1, (hsc_env1, cache1.features.enable))
 
 storeIface :: HscEnv -> ModIface -> IO ()
 storeIface _ _ =
   pure ()
 
+storeHug :: HscEnv -> Cache -> IO Cache
+storeHug hsc_env cache = do
+  pure cache {hug = Just merged}
+  where
+    merged = maybe id (unitEnv_union mergeHugs) cache.hug hsc_env.hsc_unit_env.ue_home_unit_graph
+
 finalizeCache ::
   MVar Log ->
   -- | A description of the current worker process.
@@ -630,14 +693,23 @@ finalizeCache logVar workerId hsc_env target artifacts cache0 = do
   cache1 <-
     if cache0.features.enable
     then do
-      cache1 <- fromMaybe cache0 . join <$> withHscState hsc_env \ nsNames loaderStateVar symbolCacheVar ->
-        readMVar loaderStateVar >>= traverse \ newLoaderState -> do
-          newSymbols <- readMVar symbolCacheVar
-          newNames <- readMVar nsNames
-          maybe initCache (updateCache target) cache0.interp newLoaderState (SymbolCache newSymbols) newNames cache0
+      cache1 <-
+        if cache0.features.loader
+        then do
+          fromMaybe cache0 . join <$> withHscState hsc_env \ nsNames loaderStateVar symbolCacheVar ->
+            readMVar loaderStateVar >>= traverse \ newLoaderState -> do
+              newSymbols <- readMVar symbolCacheVar
+              newNames <- readMVar nsNames
+              maybe initCache (updateCache target) cache0.interp newLoaderState (SymbolCache newSymbols) newNames cache0
+        else pure cache0
+      cache2 <-
+        if cache0.features.hpt
+        then do
+          storeHug hsc_env cache1
+        else pure cache1
       for_ artifacts \ ModuleArtifacts {iface} ->
         storeIface hsc_env iface
-      pure cache1
+      pure cache2
     else pure cache0
   report logVar workerId target cache1
   pure cache1

From 8220435def78883c111f31ae09950164c70364d9 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Mon, 7 Apr 2025 15:15:12 -0700
Subject: [PATCH 08/16] Rename loader to loaderStats

Change-Id: I4e25e45b5b19403dc15e6f7bdf3fe89f4c5d6250
---
 plugin/src/Internal/Cache.hs | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/plugin/src/Internal/Cache.hs b/plugin/src/Internal/Cache.hs
index 87fe2bee..8a1a5f2d 100644
--- a/plugin/src/Internal/Cache.hs
+++ b/plugin/src/Internal/Cache.hs
@@ -122,7 +122,7 @@ data NamesStats =
 
 data StatsUpdate =
   StatsUpdate {
-    loader :: LoaderStats,
+    loaderStats :: LoaderStats,
     symbols :: SymbolsStats,
     names :: NamesStats
   }
@@ -131,7 +131,7 @@ data StatsUpdate =
 emptyStatsUpdate :: StatsUpdate
 emptyStatsUpdate =
   StatsUpdate {
-    loader = emptyLoaderStats,
+    loaderStats = emptyLoaderStats,
     symbols = SymbolsStats {new = 0},
     names = NamesStats {new = 0}
   }
@@ -359,8 +359,8 @@ pushStats :: Bool -> Target -> Maybe LoaderStats -> SymbolsStats -> NamesStats -
 pushStats restoring target (Just new) symbols names =
   modifyStats target add
   where
-    add old | restoring = old {restore = old.restore {loader = new, symbols, names}}
-            | otherwise = old {update = old.update {loader = new, symbols, names}}
+    add old | restoring = old {restore = old.restore {loaderStats = new, symbols, names}}
+            | otherwise = old {update = old.update {loaderStats = new, symbols, names}}
 pushStats _ _ _ _ _ =
   id
 
@@ -481,18 +481,18 @@ statsMessages CacheStats {restore, update, finder} =
   hang (text "Finder:") 2 finderStats
   where
       restoreStats =
-        text (show (length restore.loader.newBcos)) <+> text "BCOs" $$
-        text (show restore.loader.linker.newClosures) <+> text "closures" $$
+        text (show (length restore.loaderStats.newBcos)) <+> text "BCOs" $$
+        text (show restore.loaderStats.linker.newClosures) <+> text "closures" $$
         text (show restore.symbols.new) <+> text "symbols" $$
-        text (show restore.loader.sameBcos) <+> text "BCOs already in cache"
+        text (show restore.loaderStats.sameBcos) <+> text "BCOs already in cache"
 
-      newBcos = text <$> update.loader.newBcos
+      newBcos = text <$> update.loaderStats.newBcos
 
       updateStats =
         (if null newBcos then text "No new BCOs" else text "New BCOs:" <+> fsep (punctuate comma newBcos)) $$
-        text (show update.loader.linker.newClosures) <+> text "new closures" $$
+        text (show update.loaderStats.linker.newClosures) <+> text "new closures" $$
         text (show update.symbols.new) <+> text "new symbols" $$
-        text (show update.loader.sameBcos) <+> text "BCOs already in cache"
+        text (show update.loaderStats.sameBcos) <+> text "BCOs already in cache"
 
       finderStats =
         hang (text "Hits:") 2 (moduleColumns finder.hits) $$

From fdff26c29118da27aa159f79cb274936c3f98f29 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Tue, 8 Apr 2025 13:55:57 -0700
Subject: [PATCH 09/16] Add haddock comments in Internal.Session

Change-Id: I0d9d2efd1dc879b8aab8fb709ec6c344bf3d98b3
---
 plugin/src/Internal/Session.hs | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/plugin/src/Internal/Session.hs b/plugin/src/Internal/Session.hs
index 6490c784..d06b5c09 100644
--- a/plugin/src/Internal/Session.hs
+++ b/plugin/src/Internal/Session.hs
@@ -42,13 +42,22 @@ import Internal.Log (Log (..), logToState)
 import Prelude hiding (log)
 import System.Environment (setEnv)
 
+-- | Worker state.
 data Env =
   Env {
+    -- | Logger used to receive messages from GHC and relay them to Buck.
     log :: MVar Log,
+
+    -- | Parts of @HscEnv@ we share between sessions.
     cache :: MVar Cache,
+
+    -- | Preprocessed command line args from Buck.
     args :: Args
   }
 
+-- | Add all the directories passed by Buck in @--bin-path@ options to the global @$PATH@.
+-- Although Buck intends these to be module specific, all subsequent compile jobs will see all previous jobs' entries,
+-- since we only have one process environment.
 setupPath :: Args -> Cache -> IO Cache
 setupPath args old = do
   setEnv "PATH" (intercalate ":" (toList path.extra ++ maybeToList path.initial))
@@ -68,6 +77,8 @@ setTempDir dir = hscUpdateFlags \ dflags -> dflags {tmpDir = TempDir dir}
 dummyLocation :: a -> Located a
 dummyLocation = mkGeneralLocated "by Buck2"
 
+-- | Parse command line flags into @DynFlags@ and set up the logger. Extracted from GHC.
+-- Returns the subset of args that have not been recognized as options.
 parseFlags :: [Located String] -> Ghc (DynFlags, Logger, [Located String], DriverMessages)
 parseFlags argv = do
   dflags0 <- GHC.getSessionDynFlags
@@ -77,6 +88,8 @@ parseFlags argv = do
   (dflags, fileish_args, dynamicFlagWarnings) <- parseDynamicFlags logger2 dflags1 argv
   pure (dflags, setLogFlags logger2 (initLogFlags dflags), fileish_args, dynamicFlagWarnings)
 
+-- | Parse CLI args and set up the GHC session.
+-- Returns the subset of args that have not been recognized as options.
 initGhc ::
   DynFlags ->
   Logger ->
@@ -92,6 +105,10 @@ initGhc dflags0 logger fileish_args dynamicFlagWarnings = do
   where
     flagWarnings' = GhcDriverMessage <$> dynamicFlagWarnings
 
+-- | Run a program with a fresh session constructed from command line args.
+-- Passes the unprocessed args to the callback, which usually consist of the file or module names intended for
+-- compilation.
+-- In a Buck compile step these should always be a single path, but in the metadata step they enumerate an entire unit.
 withGhcInSession :: Env -> ([(String, Maybe Phase)] -> Ghc a) -> [Located String] -> Ghc a
 withGhcInSession env prog argv = do
   pushLogHookM (const (logToState env.log))
@@ -100,6 +117,9 @@ withGhcInSession env prog argv = do
     srcs <- initGhc dflags0 logger fileish_args dynamicFlagWarnings
     prog srcs
 
+-- | Create a base session and store it in the cache.
+-- On subsequent calls, return the cached session, unless the cache is disabled or @reuse@ is true.
+-- This will at some point be replaced by more deliberate methods.
 ensureSession :: MVar Cache -> Args -> IO HscEnv
 ensureSession cacheVar args =
   modifyMVar cacheVar \ cache -> do
@@ -108,6 +128,8 @@ ensureSession cacheVar args =
     then pure (cache {baseSession = Just newEnv}, newEnv)
     else pure (cache, newEnv)
 
+-- | Run a @Ghc@ program to completion with a fresh clone of the base session.
+-- See 'ensureSession' for @reuse@.
 runSession :: Env -> ([Located String] -> Ghc (Maybe a)) -> IO (Maybe a)
 runSession Env {log, args, cache} prog = do
   modifyMVar_ cache (setupPath args)
@@ -117,12 +139,16 @@ runSession Env {log, args, cache} prog = do
     traverse_ (modifySession . setTempDir) args.tempDir
     handleExceptions log Nothing (prog (map dummyLocation args.ghcOptions))
 
+-- | When compiling a module, the leftover arguments from parsing @DynFlags@ should be a single source file path.
+-- Wrap it in 'Target' or terminate.
 ensureSingleTarget :: [(String, Maybe Phase)] -> Ghc Target
 ensureSingleTarget = \case
   [(src, Nothing)] -> pure (Target src)
   [(_, phase)] -> panic ("Called worker with unexpected start phase: " ++ show phase)
   args -> panic ("Called worker with multiple targets: " ++ show args)
 
+-- | Run a @Ghc@ program to completion with a fresh clone of the base session, wrapped in a handler operating on a
+-- compilation target.
 withGhcUsingCache :: (Target -> Ghc a -> Ghc (Maybe b)) -> Env -> (Target -> Ghc a) -> IO (Maybe b)
 withGhcUsingCache cacheHandler env prog =
   runSession env $ withGhcInSession env \ srcs -> do
@@ -131,6 +157,8 @@ withGhcUsingCache cacheHandler env prog =
       initializeSessionPlugins
       prog target
 
+-- | Run a @Ghc@ program to completion with a fresh clone of the base session augmented by some persisted state.
+-- This is a compat shim for the multiplex worker.
 withGhc :: Env -> (Target -> Ghc (Maybe a)) -> IO (Maybe a)
 withGhc env =
   withGhcUsingCache cacheHandler env
@@ -143,6 +171,8 @@ withGhc env =
           pure (Nothing, a)
       pure (snd <$> result)
 
+-- | Run a @Ghc@ program to completion with a fresh clone of the base session augmented by some persisted state.
+-- Return the interface and bytecode.
 withGhcDefault :: Env -> (Target -> Ghc (Maybe (Maybe ModuleArtifacts, a))) -> IO (Maybe (Maybe ModuleArtifacts, a))
 withGhcDefault env =
   withGhcUsingCache (withCache env.log env.args.workerTargetId env.cache) env

From 834382d96b860c8c0d8541300f61b60b73ed0994 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Tue, 8 Apr 2025 14:01:13 -0700
Subject: [PATCH 10/16] withGhcMhu: run GHC session with Multi-Home-Unit Cache
 support

Change-Id: Iccb0512ff1c262b1b53e20e089a14e705cd0cc6e
---
 plugin/src/Internal/Session.hs | 130 ++++++++++++++++++++++++++++++---
 1 file changed, 118 insertions(+), 12 deletions(-)

diff --git a/plugin/src/Internal/Session.hs b/plugin/src/Internal/Session.hs
index d06b5c09..803c2d0c 100644
--- a/plugin/src/Internal/Session.hs
+++ b/plugin/src/Internal/Session.hs
@@ -5,7 +5,7 @@ import Control.Monad (unless)
 import Control.Monad.IO.Class (liftIO)
 import Data.Foldable (toList, traverse_)
 import Data.IORef (newIORef)
-import Data.List (intercalate)
+import Data.List (intercalate, isPrefixOf)
 import Data.List.NonEmpty (nonEmpty)
 import Data.Maybe (maybeToList)
 import qualified Data.Set as Set
@@ -82,7 +82,7 @@ dummyLocation = mkGeneralLocated "by Buck2"
 parseFlags :: [Located String] -> Ghc (DynFlags, Logger, [Located String], DriverMessages)
 parseFlags argv = do
   dflags0 <- GHC.getSessionDynFlags
-  let dflags1 = dflags0 {ghcMode = OneShot, ghcLink = LinkBinary, verbosity = 0}
+  let dflags1 = dflags0 {ghcLink = LinkBinary, verbosity = 0}
   logger1 <- getLogger
   let logger2 = setLogFlags logger1 (initLogFlags dflags1)
   (dflags, fileish_args, dynamicFlagWarnings) <- parseDynamicFlags logger2 dflags1 argv
@@ -120,20 +120,23 @@ withGhcInSession env prog argv = do
 -- | Create a base session and store it in the cache.
 -- On subsequent calls, return the cached session, unless the cache is disabled or @reuse@ is true.
 -- This will at some point be replaced by more deliberate methods.
-ensureSession :: MVar Cache -> Args -> IO HscEnv
-ensureSession cacheVar args =
+ensureSession :: Bool -> MVar Cache -> Args -> IO HscEnv
+ensureSession reuse cacheVar args =
   modifyMVar cacheVar \ cache -> do
-    newEnv <- maybe (initHscEnv args.topdir) pure cache.baseSession
-    if cache.features.enable
-    then pure (cache {baseSession = Just newEnv}, newEnv)
-    else pure (cache, newEnv)
+    if cache.features.enable && reuse
+    then do
+      newEnv <- maybe (initHscEnv args.topdir) pure cache.baseSession
+      pure (cache {baseSession = Just newEnv}, newEnv)
+    else do
+      newEnv <- initHscEnv args.topdir
+      pure (cache, newEnv)
 
 -- | Run a @Ghc@ program to completion with a fresh clone of the base session.
 -- See 'ensureSession' for @reuse@.
-runSession :: Env -> ([Located String] -> Ghc (Maybe a)) -> IO (Maybe a)
-runSession Env {log, args, cache} prog = do
+runSession :: Bool -> Env -> ([Located String] -> Ghc (Maybe a)) -> IO (Maybe a)
+runSession reuse Env {log, args, cache} prog = do
   modifyMVar_ cache (setupPath args)
-  hsc_env <- ensureSession cache args
+  hsc_env <- ensureSession reuse cache args
   session <- Session <$> newIORef hsc_env
   flip unGhc session $ withSignalHandlers do
     traverse_ (modifySession . setTempDir) args.tempDir
@@ -151,7 +154,7 @@ ensureSingleTarget = \case
 -- compilation target.
 withGhcUsingCache :: (Target -> Ghc a -> Ghc (Maybe b)) -> Env -> (Target -> Ghc a) -> IO (Maybe b)
 withGhcUsingCache cacheHandler env prog =
-  runSession env $ withGhcInSession env \ srcs -> do
+  runSession True env $ withGhcInSession env \ srcs -> do
     target <- ensureSingleTarget srcs
     cacheHandler target do
       initializeSessionPlugins
@@ -176,3 +179,106 @@ withGhc env =
 withGhcDefault :: Env -> (Target -> Ghc (Maybe (Maybe ModuleArtifacts, a))) -> IO (Maybe (Maybe ModuleArtifacts, a))
 withGhcDefault env =
   withGhcUsingCache (withCache env.log env.args.workerTargetId env.cache) env
+
+
+-- | Command line args that have to be stored in the current home unit env.
+-- These are specified as a single program argument with their option argument, without whitespace in between.
+specificPrefixSwitches :: [String]
+specificPrefixSwitches =
+  [
+    "-i"
+  ]
+
+-- | Command line args that have to be stored in the current home unit env.
+specificSwitches :: [String]
+specificSwitches =
+  [
+    "-o",
+    "-dyno",
+    "-ohi",
+    "-dynohi",
+    "-this-unit-id",
+    "-package",
+    "-package-id",
+    "-stubdir"
+  ]
+
+-- | Indicate whether the CLI arg starts with any of the values in 'specificPrefixSwitches'.
+isSpecificPrefix :: String -> Bool
+isSpecificPrefix arg =
+  any (`isPrefixOf` arg) specificPrefixSwitches
+
+-- | Indicate whether the CLI arg is in 'specificSwitches'.
+isSpecific :: String -> Bool
+isSpecific =
+  flip elem specificSwitches
+
+-- | Separate the command line given by Buck into options pertaining to the target home unit and the rest.
+-- Write the rest back to the 'Env' passed to the continuation for processing as global args, and pass the home unit
+-- specific args as the second argument to the continuation.
+--
+-- @-hide-all-packages@ is removed entirely, which may be obsolete.
+-- @-this-unit-id@ is added to both parts, since the global session is always initialized with a default session.
+withUnitSpecificOptions :: Bool -> Env -> (Env -> [String] -> [Located String] -> Ghc (Maybe a)) -> IO (Maybe a)
+withUnitSpecificOptions reuse env use =
+  runSession reuse env1 $ use env1 specific
+  where
+    env1 = env {args = env.args {ghcOptions = general}}
+    (general, specific) = spin ([], []) env.args.ghcOptions
+
+    spin (g, s) = \case
+      [] -> (reverse g, reverse s)
+      "-hide-all-packages" : rest
+        -> spin (g, s) rest
+      "-this-unit-id" : uid : rest
+        -> spin (uid : "-this-unit-id" : g, uid : "-this-unit-id" : s) rest
+      switch : arg : rest
+        | isSpecific switch
+        -> spin (g, arg : switch : s) rest
+      arg : rest
+        | isSpecificPrefix arg
+        -> spin (g, arg : s) rest
+        | otherwise
+        -> spin (arg : g, s) rest
+
+-- | Run a GHC session with multiple home unit support, separating the CLI args for the current unit from the rest.
+withGhcInSessionMhu ::
+  Env ->
+  ([String] -> [(String, Maybe Phase)] -> Ghc (Maybe a)) ->
+  IO (Maybe a)
+withGhcInSessionMhu env prog =
+  withUnitSpecificOptions True env \ env1 specific -> withGhcInSession env1 (prog specific)
+
+-- | Like @withGhcInSessionMhu@, but wrap with the given function operating on the current target for caching purposes.
+withGhcUsingCacheMhu ::
+  (Target -> Ghc a -> Ghc (Maybe b)) ->
+  Env ->
+  ([String] -> Target -> Ghc a) ->
+  IO (Maybe b)
+withGhcUsingCacheMhu cacheHandler env prog =
+  withGhcInSessionMhu env \ specific srcs -> do
+    target <- ensureSingleTarget srcs
+    cacheHandler target do
+      initializeSessionPlugins
+      prog specific target
+
+-- | Like @withGhcUsingCacheMhu@, using the default cache handler @withCache@.
+withGhcMhu :: Env -> ([String] -> Target -> Ghc (Maybe a)) -> IO (Maybe a)
+withGhcMhu env =
+  withGhcUsingCacheMhu cacheHandler env
+  where
+    cacheHandler target prog = do
+      result <- withCache env.log env.args.workerTargetId env.cache target do
+        res <- prog
+        pure do
+          a <- res
+          pure (Nothing, a)
+      pure (snd <$> result)
+
+-- | Like @withGhcMhu@, specialized to @ModuleArtifacts@.
+withGhcMhuDefault ::
+  Env ->
+  ([String] -> Target -> Ghc (Maybe (Maybe ModuleArtifacts, a))) ->
+  IO (Maybe (Maybe ModuleArtifacts, a))
+withGhcMhuDefault env =
+  withGhcUsingCacheMhu (withCache env.log env.args.workerTargetId env.cache) env

From c8603bb9b75d04ea249901483822cb4a4b057149 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Tue, 8 Apr 2025 13:30:55 -0700
Subject: [PATCH 11/16] Introduce new Internal.CompileHpt module and
 compileModuleWithDepsInHpt.

rename original compile with compileModuleWithDepsInEps

Change-Id: I45af5159cdf3beec5a835ca7e607563f0ad4d965
---
 buck-worker/Main.hs                       |  68 +++++++---
 plugin/ghc-persistent-worker-plugin.cabal |   1 +
 plugin/src/BUCK                           |   1 +
 plugin/src/GHCPersistentWorkerPlugin.hs   |   4 +-
 plugin/src/Internal/Compile.hs            |   4 +-
 plugin/src/Internal/CompileHpt.hs         | 144 ++++++++++++++++++++++
 plugin/src/Internal/Error.hs              |  15 ++-
 7 files changed, 215 insertions(+), 22 deletions(-)
 create mode 100644 plugin/src/Internal/CompileHpt.hs

diff --git a/buck-worker/Main.hs b/buck-worker/Main.hs
index 665cb9ec..325c5524 100644
--- a/buck-worker/Main.hs
+++ b/buck-worker/Main.hs
@@ -2,7 +2,8 @@
 
 module Main where
 
-import BuckArgs (BuckArgs (abiOut), CompileResult (..), parseBuckArgs, toGhcArgs, writeResult)
+import BuckArgs (BuckArgs (abiOut), CompileResult (..), Mode (..), parseBuckArgs, toGhcArgs, writeResult)
+import qualified BuckArgs (BuckArgs (mode))
 import BuckWorker (ExecuteCommand, ExecuteCommand'EnvironmentEntry, ExecuteEvent, ExecuteResponse, Worker (..), Instrument (..))
 import Control.Concurrent (MVar, modifyMVar_, newMVar)
 import Control.Concurrent.Async (async, cancel, wait)
@@ -19,35 +20,48 @@ import Control.Exception (
 import Control.Monad (foldM, void, when, forever)
 import Control.Monad.IO.Class (liftIO)
 import Data.Foldable (for_)
+import Data.Functor ((<&>))
+import Data.Int (Int32)
 import Data.List (dropWhileEnd)
 import Data.Map (Map)
 import Data.Map.Strict qualified as Map
 import Data.Text qualified as Text
 import Data.Text.Encoding (decodeUtf8Lenient)
+import GHC (DynFlags (..), Ghc, getSession)
+import GHC.Driver.DynFlags (GhcMode (..))
+import GHC.Driver.Env (hscUpdateFlags)
+import GHC.Driver.Monad (modifySession)
 import GHC (Ghc, getSession)
 import GHC.IO.Handle.FD (withFileBlocking)
-import GHC.Stats (getRTSStats, RTSStats (..), GCDetails (..))
+import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
 import Internal.AbiHash (AbiHash (..), showAbiHash)
 import Internal.Cache (Cache (..), CacheFeatures (..), ModuleArtifacts (..), Target (..), emptyCache, emptyCacheWith)
-import Internal.Compile (compile)
+import Internal.Compile (compileModuleWithDepsInEps)
+import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (dbg, logFlush, newLog)
-import Internal.Session (Env (..), withGhc)
+import Internal.Session (Env (..), withGhcMhu)
 import Network.GRPC.Client (Connection, Server (ServerUnix), recvNextOutput, sendFinalInput, withConnection, withRPC)
 import Network.GRPC.Common (NextElem (..), Proxy (..), def)
 import Network.GRPC.Common.Protobuf (Proto, Protobuf, defMessage, (%~), (&), (.~), (^.))
 import Network.GRPC.Server.Protobuf (ProtobufMethodsOf)
-import Network.GRPC.Server.Run (InsecureConfig (InsecureUnix), ServerConfig (..), runServerWithHandlers)
-import Network.GRPC.Server.StreamType (Methods (Method, NoMoreMethods), fromMethods, mkClientStreaming, mkNonStreaming, mkServerStreaming, simpleMethods)
+import Network.GRPC.Server.Run (InsecureConfig (..), ServerConfig (..), runServerWithHandlers)
+import Network.GRPC.Server.StreamType (
+  Methods (..),
+  fromMethods,
+  mkClientStreaming,
+  mkNonStreaming,
+  mkServerStreaming,
+  simpleMethods,
+  )
 import Prelude hiding (log)
+import qualified Proto.Instrument as Instr
+import Proto.Instrument_Fields qualified as Instr
+import Proto.Worker_Fields qualified as Fields
 import System.Directory (createDirectory, removeFile)
 import System.Environment (getArgs, getEnv)
 import System.Exit (exitFailure)
 import System.FilePath (takeDirectory, (</>))
-import System.IO (BufferMode (LineBuffering), IOMode (..), hGetLine, hPutStrLn, hSetBuffering, stderr, stdout, hPutStr)
-
-import Proto.Worker_Fields qualified as Fields
-import Proto.Instrument_Fields qualified as Instr
-import qualified Proto.Instrument as Instr
+import System.IO (BufferMode (..), IOMode (..), hGetLine, hPutStr, hPutStrLn, hSetBuffering, stderr, stdout)
 
 data WorkerStatus =
   WorkerStatus {
@@ -74,15 +88,17 @@ compileAndReadAbiHash ::
   WorkerMode ->
   Chan (Proto Instr.Event) ->
   BuckArgs ->
+  [String] ->
   Target ->
   Ghc (Maybe CompileResult)
-compileAndReadAbiHash _mode instrChan args target = do
+compileAndReadAbiHash mode instrChan args specific target = do
 
   liftIO $ writeChan instrChan $
     defMessage &
       Instr.compileStart .~
         compileStart target.get
 
+  modifySession $ hscUpdateFlags \ d -> d {ghcMode}
   compile target >>= traverse \ artifacts -> do
     hsc_env <- getSession
     let
@@ -91,6 +107,27 @@ compileAndReadAbiHash _mode instrChan args target = do
         path <- args.abiOut
         Just AbiHash {path, hash = showAbiHash hsc_env artifacts.iface}
     pure CompileResult {artifacts, abiHash}
+  where
+    (ghcMode, compile) = case mode of
+      WorkerOneshotMode -> (OneShot, compileModuleWithDepsInEps)
+      WorkerMakeMode -> (CompManager, compileModuleWithDepsInHpt specific)
+
+-- | Process a worker request based on the operational mode specified in the request arguments, either compiling a
+-- single module for 'ModeCompile' (@-c@), or computing and writing the module graph to a JSON file for 'ModeMetadata'
+-- (@-M@).
+dispatch ::
+  WorkerMode ->
+  Chan (Proto Instr.Event) ->
+  Env ->
+  BuckArgs ->
+  IO Int32
+dispatch workerMode instrChan env args =
+  case args.mode of
+    Just ModeCompile -> do
+      result <- withGhcMhu env (compileAndReadAbiHash workerMode instrChan args)
+      writeResult args result
+    Just m -> error ("worker: mode not implemented: " ++ show m)
+    Nothing -> error "worker: no mode specified"
 
 compileStart :: String -> Proto Instr.CompileStart
 compileStart target =
@@ -168,11 +205,10 @@ execute status cache mode instrChan req = do
       bracket (startJob status) (finishJob status) \ _ -> do
         log <- newLog True
         let env = Env {log, cache, args}
-        result <- withGhc env (compileAndReadAbiHash mode instrChan buckArgs)
-        pure (env, buckArgs, result)
+        result <- dispatch mode instrChan env buckArgs
+        pure (env, result)
 
-    successResponse (env, buckArgs, result) = do
-      exitCode <- writeResult buckArgs result
+    successResponse (env, exitCode) = do
       output <- logFlush env.log
       pure $
         defMessage
diff --git a/plugin/ghc-persistent-worker-plugin.cabal b/plugin/ghc-persistent-worker-plugin.cabal
index 78a1955c..2359ca47 100644
--- a/plugin/ghc-persistent-worker-plugin.cabal
+++ b/plugin/ghc-persistent-worker-plugin.cabal
@@ -29,6 +29,7 @@ library
                       Internal.Args
                       Internal.Cache
                       Internal.Compile
+                      Internal.CompileHpt
                       Internal.Error
                       Internal.Log
                       Internal.Session
diff --git a/plugin/src/BUCK b/plugin/src/BUCK
index 982cf9d4..1f1b7700 100644
--- a/plugin/src/BUCK
+++ b/plugin/src/BUCK
@@ -3,6 +3,7 @@
     "Args",
     "Cache",
     "Compile",
+    "CompileHpt",
     "Error",
     "Log",
     "Session",
diff --git a/plugin/src/GHCPersistentWorkerPlugin.hs b/plugin/src/GHCPersistentWorkerPlugin.hs
index 7e0a3a6b..e348768c 100644
--- a/plugin/src/GHCPersistentWorkerPlugin.hs
+++ b/plugin/src/GHCPersistentWorkerPlugin.hs
@@ -21,7 +21,7 @@ import GHC.Runtime.Interpreter.Types (Interp (..), InterpInstance (..))
 import GHC.Settings.Config (cProjectVersion)
 import Internal.Args (Args (..), emptyArgs)
 import Internal.Cache (emptyCache)
-import Internal.Compile (compile)
+import Internal.Compile (compileModuleWithDepsInEps)
 import Internal.Error (handleExceptions)
 import Internal.Log (Log, logFlushBytes, logToState, newLog)
 import Internal.Session (Env (..), withGhc)
@@ -160,7 +160,7 @@ workerImplCustom =
     use cache logVar args = do
       let env' = Env {log = logVar, cache, args = (emptyArgs mempty) {ghcOptions = sanitize args}}
       fmap (fromMaybe 2) $ withGhc env' \ target ->
-        compile target >>= \case
+        compileModuleWithDepsInEps target >>= \case
           Just _ -> pure (Just 0)
           Nothing -> pure (Just 1)
 
diff --git a/plugin/src/Internal/Compile.hs b/plugin/src/Internal/Compile.hs
index dca56cfb..df25373e 100644
--- a/plugin/src/Internal/Compile.hs
+++ b/plugin/src/Internal/Compile.hs
@@ -91,7 +91,7 @@ compileFile hsc_env src = do
     pipe_env = mkPipeEnv NoStop offset_file Nothing output
     pipeline = pipelineOneshot pipe_env (setDumpPrefix pipe_env hsc_env) offset_file
 
-compile :: Target -> Ghc (Maybe ModuleArtifacts)
-compile (Target src) = do
+compileModuleWithDepsInEps :: Target -> Ghc (Maybe ModuleArtifacts)
+compileModuleWithDepsInEps (Target src) = do
   hsc_env <- liftIO . initializePlugins =<< getSession
   liftIO $ compileFile hsc_env src
diff --git a/plugin/src/Internal/CompileHpt.hs b/plugin/src/Internal/CompileHpt.hs
new file mode 100644
index 00000000..1c27307e
--- /dev/null
+++ b/plugin/src/Internal/CompileHpt.hs
@@ -0,0 +1,144 @@
+{-# LANGUAGE ViewPatterns, CPP, OverloadedStrings #-}
+
+module Internal.CompileHpt where
+
+import Control.Monad (when)
+import Data.Foldable (fold)
+import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
+import Data.Set (Set)
+import GHC (
+  DynFlags (..),
+  GeneralFlag (..),
+  Ghc,
+  GhcMonad (..),
+  Logger,
+  ModLocation (..),
+  ModSummary (..),
+  gopt,
+  mkGeneralLocated,
+  setUnitDynFlags,
+  )
+import GHC.Driver.Config.Diagnostic (initDiagOpts, initPrintConfig)
+import GHC.Driver.Env (HscEnv (..), hscSetActiveUnitId, hscUpdateHUG)
+import GHC.Driver.Errors (printOrThrowDiagnostics)
+import GHC.Driver.Errors.Types (GhcMessage (..))
+import GHC.Driver.Make (summariseFile)
+import GHC.Driver.Monad (modifySession)
+import GHC.Driver.Pipeline (compileOne)
+import GHC.Driver.Session (parseDynamicFlagsCmdLine)
+import GHC.Runtime.Loader (initializeSessionPlugins)
+import GHC.Unit (UnitId, UnitState (..), stringToUnitId, unitIdString)
+import GHC.Unit.Env (
+  HomeUnitEnv (..),
+  HomeUnitGraph,
+  UnitEnv (..),
+  UnitEnvGraph (..),
+  addHomeModInfoToHug,
+  ue_unsafeHomeUnit,
+  unitEnv_lookup_maybe,
+  )
+import GHC.Unit.Home.ModInfo (HomeModInfo (..), HomeModLinkable (..))
+import GHC.Utils.Monad (MonadIO (..))
+import GHC.Utils.TmpFs (TmpFs, cleanCurrentModuleTempFiles, keepCurrentModuleTempFiles)
+import Internal.Cache (ModuleArtifacts (..), Target (..))
+import Internal.Error (eitherMessages)
+
+-- | Insert a compilation result into the current unit's home package table, as it is done by upsweep.
+addDepsToHscEnv :: [HomeModInfo] -> HscEnv -> HscEnv
+addDepsToHscEnv deps = hscUpdateHUG (\hug -> foldr addHomeModInfoToHug hug deps)
+
+-- | Update the location of the result of @summariseFile@ to point to the locations specified on the command line, since
+-- these are placed in the source file's directory by that function.
+setHiLocation :: HscEnv -> ModSummary -> ModSummary
+setHiLocation HscEnv {hsc_dflags = DynFlags {outputHi = Just ml_hi_file, outputFile_ = Just ml_obj_file}} summ =
+  summ {ms_location = summ.ms_location {ml_hi_file, ml_obj_file}}
+setHiLocation _ summ = summ
+
+-- | Parse command line flags, used to create unit-specific @DynFlags@.
+unitFlags :: [String] -> HscEnv -> Ghc DynFlags
+unitFlags args HscEnv {hsc_logger, hsc_dflags = dflags0} = do
+  (dflags, _, warns) <- parseDynamicFlagsCmdLine dflags0 (map (mkGeneralLocated "no loc") args)
+  liftIO $ printOrThrowDiagnostics hsc_logger (initPrintConfig dflags) (initDiagOpts dflags) (GhcDriverMessage <$> warns)
+  pure dflags
+
+-- | Turn each @-package@ argument in the given list into a @-package-id@ argument if the following package name refers
+-- to a home unit, and create a set of all mentioned home unit names.
+-- This is needed for GHC to recognize home unit dependencies.
+adaptHp :: HomeUnitGraph -> [String] -> (Set UnitId, [String])
+adaptHp (UnitEnvGraph ueg) =
+  spin mempty
+  where
+    spin seen = \case
+      "-package" : p : rest
+        | Map.member (stringToUnitId p) ueg
+        -> (["-package-id", p] ++) <$> spin (Set.insert (stringToUnitId p) seen) rest
+      arg : rest -> (arg :) <$> spin seen rest
+      [] -> (seen, [])
+
+-- | Return the given unit's dependencies.
+homeUnitDeps :: HscEnv -> UnitId -> Maybe [UnitId]
+homeUnitDeps hsc_env target = do
+  HomeUnitEnv {homeUnitEnv_units = UnitState {homeUnitDepends}} <- unitEnv_lookup_maybe target hug
+  pure homeUnitDepends
+  where
+    hug = hsc_env.hsc_unit_env.ue_home_unit_graph
+
+-- | Assemble @-package-id@ arguments for the current unit's dependencies, omitting those that are present in the
+-- provided set, which are the current module's deps specified by Buck.
+homeUnitDepFlags :: HscEnv -> Set UnitId -> UnitId -> [String]
+homeUnitDepFlags hsc_env explicit target =
+  concat [["-package-id", unitIdString u] | u <- fold prev_deps, not (Set.member u explicit)]
+  where
+    prev_deps = homeUnitDeps hsc_env target
+
+-- | Update the current unit's @DynFlags@ stored in the unit env, and reinitialize its unit state.
+-- Since different modules in the same unit may have arbitrary subsets of the unit's package dependencies when Buck
+-- compiles them, we take the union of existing and new dependencies.
+initUnit :: [String] -> Ghc ()
+initUnit specific = do
+  hsc_env0 <- getSession
+  let current = hsc_env0.hsc_unit_env.ue_current_unit
+      (explicit, withPackageId) = adaptHp hsc_env0.hsc_unit_env.ue_home_unit_graph specific
+      unitOptions = withPackageId ++ homeUnitDepFlags hsc_env0 explicit current
+  dflags <- unitFlags unitOptions hsc_env0
+  setUnitDynFlags current dflags
+  modifySession (hscSetActiveUnitId current)
+
+-- | Not used yet.
+cleanCurrentModuleTempFilesMaybe :: MonadIO m => Logger -> TmpFs -> DynFlags -> m ()
+cleanCurrentModuleTempFilesMaybe logger tmpfs dflags =
+  if gopt Opt_KeepTmpFiles dflags
+    then liftIO $ keepCurrentModuleTempFiles logger tmpfs
+    else liftIO $ cleanCurrentModuleTempFiles logger tmpfs
+
+-- | Compile a module with multiple home units in the session state, using the home package table to look up
+-- dependencies.
+--
+-- First, update the current unit's configuration to include this module's dependencies.
+-- Buck only provides @-package@ flags for deps that are used by a given module, while the unit state is designed to be
+-- initialized up front with the deps of all modules.
+-- Note: This should soon be obsolete, since we now have full control over the metadata step.
+--
+-- Next, perform the steps that usually happen in make mode's upsweep:
+-- - Create a @ModSummary@ using @summariseFile@
+-- - Call the module compilation function @compileOne@
+-- - Store the resulting @HomeModInfo@ in the current unit's home package table.
+compileModuleWithDepsInHpt ::
+  [String] ->
+  Target ->
+  Ghc (Maybe ModuleArtifacts)
+compileModuleWithDepsInHpt specific (Target src) = do
+  initializeSessionPlugins
+  initUnit specific
+  hsc_env <- getSession
+  hmi@HomeModInfo {hm_iface = iface, hm_linkable} <- liftIO do
+    summResult <- summariseFile hsc_env (ue_unsafeHomeUnit (hsc_unit_env hsc_env)) mempty src Nothing Nothing
+    summary <- setHiLocation hsc_env <$> eitherMessages GhcDriverMessage summResult
+    result <- compileOne hsc_env summary 1 100000 Nothing (HomeModLinkable Nothing Nothing)
+    -- This deletes assembly files too early
+    when False do
+      cleanCurrentModuleTempFilesMaybe (hsc_logger hsc_env) (hsc_tmpfs hsc_env) summary.ms_hspp_opts
+    pure result
+  modifySession (addDepsToHscEnv [hmi])
+  pure (Just ModuleArtifacts {iface, bytecode = homeMod_bytecode hm_linkable})
diff --git a/plugin/src/Internal/Error.hs b/plugin/src/Internal/Error.hs
index 2ef60927..e03078b3 100644
--- a/plugin/src/Internal/Error.hs
+++ b/plugin/src/Internal/Error.hs
@@ -5,12 +5,14 @@ module Internal.Error where
 import Control.Concurrent.MVar (MVar)
 import Control.Exception (AsyncException (..), Exception (..), IOException, throwIO)
 import qualified Control.Monad.Catch as MC
-import Control.Monad.IO.Class (liftIO)
+import Control.Monad.IO.Class (liftIO, MonadIO)
 import GHC (Ghc, GhcException (..), printException)
-import GHC.Types.SourceError (SourceError)
+import GHC.Types.SourceError (SourceError, throwErrors)
 import Internal.Log (Log, logOther)
 import System.Environment (getProgName)
 import System.Exit (ExitCode)
+import GHC.Driver.Errors.Types (GhcMessage)
+import GHC.Types.Error (Messages)
 
 handleExceptions :: MVar Log -> a -> Ghc a -> Ghc a
 handleExceptions logVar errResult =
@@ -47,3 +49,12 @@ handleExceptions logVar errResult =
       = fm (show (Panic (show exception)))
 
     fm = logOther logVar
+
+eitherMessages ::
+  MonadIO m =>
+  (a -> GhcMessage) ->
+  Either (Messages a) b ->
+  m b
+eitherMessages toMessage = \case
+  Right b -> pure b
+  Left errs -> throwErrors (toMessage <$> errs)

From ef36296a199209af11ecbbcde803de03bc0f086c Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Tue, 8 Apr 2025 14:16:34 -0700
Subject: [PATCH 12/16] Internalize computeMetadata.

Inlined GHC.Driver.MakeFile, GHC.Driver.MakeFile.JSON and then implemented
computeMetadata internally here.

Change-Id: Id713433734861d26f889e48e5401c716e90f66c1
---
 buck-worker/Main.hs                       |   1 +
 plugin/ghc-persistent-worker-plugin.cabal |   3 +
 plugin/src/BUCK                           |   3 +
 plugin/src/Internal/MakeFile.hs           | 536 ++++++++++++++++++++++
 plugin/src/Internal/MakeFile/JSON.hs      | 240 ++++++++++
 plugin/src/Internal/Metadata.hs           |  53 +++
 6 files changed, 836 insertions(+)
 create mode 100644 plugin/src/Internal/MakeFile.hs
 create mode 100644 plugin/src/Internal/MakeFile/JSON.hs
 create mode 100644 plugin/src/Internal/Metadata.hs

diff --git a/buck-worker/Main.hs b/buck-worker/Main.hs
index 325c5524..60dfe567 100644
--- a/buck-worker/Main.hs
+++ b/buck-worker/Main.hs
@@ -39,6 +39,7 @@ import Internal.Cache (Cache (..), CacheFeatures (..), ModuleArtifacts (..), Tar
 import Internal.Compile (compileModuleWithDepsInEps)
 import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (dbg, logFlush, newLog)
+import Internal.Metadata (computeMetadata)
 import Internal.Session (Env (..), withGhcMhu)
 import Network.GRPC.Client (Connection, Server (ServerUnix), recvNextOutput, sendFinalInput, withConnection, withRPC)
 import Network.GRPC.Common (NextElem (..), Proxy (..), def)
diff --git a/plugin/ghc-persistent-worker-plugin.cabal b/plugin/ghc-persistent-worker-plugin.cabal
index 2359ca47..7fe19c77 100644
--- a/plugin/ghc-persistent-worker-plugin.cabal
+++ b/plugin/ghc-persistent-worker-plugin.cabal
@@ -32,6 +32,9 @@ library
                       Internal.CompileHpt
                       Internal.Error
                       Internal.Log
+                      Internal.MakeFile
+                      Internal.MakeFile.JSON
+                      Internal.Metadata
                       Internal.Session
                       Util
     other-modules:  GHC.Main
diff --git a/plugin/src/BUCK b/plugin/src/BUCK
index 1f1b7700..fc104550 100644
--- a/plugin/src/BUCK
+++ b/plugin/src/BUCK
@@ -6,5 +6,8 @@
     "CompileHpt",
     "Error",
     "Log",
+    "MakeFile",
+    "MakeFile/JSON",
+    "Metadata",
     "Session",
 ]]
diff --git a/plugin/src/Internal/MakeFile.hs b/plugin/src/Internal/MakeFile.hs
new file mode 100644
index 00000000..928d850b
--- /dev/null
+++ b/plugin/src/Internal/MakeFile.hs
@@ -0,0 +1,536 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# language NoImplicitPrelude, FieldSelectors, CPP #-}
+{-# OPTIONS_GHC -Wno-name-shadowing #-}
+
+-----------------------------------------------------------------------------
+--
+-- Makefile Dependency Generation
+--
+-- (c) The University of Glasgow 2005
+--
+-----------------------------------------------------------------------------
+
+module Internal.MakeFile where
+
+import GHC.Prelude
+
+import qualified GHC
+import GHC.Data.Maybe
+import GHC.Driver.Monad
+import GHC.Driver.DynFlags
+import GHC.Driver.Ppr
+import Internal.MakeFile.JSON
+import GHC.Utils.Misc
+import GHC.Driver.Env
+import GHC.Driver.Errors.Types
+import GHC.Driver.Pipeline (runPipeline, TPhase (T_Unlit, T_FileArgs), use, mkPipeEnv)
+import GHC.Driver.Phases (StopPhase (StopPreprocess), startPhase, Phase (Unlit))
+import GHC.Driver.Pipeline.Monad (PipelineOutput (NoOutputFile))
+import GHC.Driver.Session (pgm_F)
+import qualified GHC.SysTools as SysTools
+import GHC.Data.Graph.Directed ( SCC(..) )
+import GHC.Utils.Outputable
+import GHC.Utils.Panic
+import GHC.Types.SourceError
+import GHC.Types.SrcLoc
+import GHC.Types.PkgQual
+import Data.List (partition)
+import GHC.Utils.TmpFs
+
+import GHC.Iface.Load (cannotFindModule)
+import GHC.Iface.Errors.Types
+
+import GHC.Unit.Module
+import GHC.Unit.Module.ModSummary
+import GHC.Unit.Module.Graph
+import GHC.Unit.Finder
+import GHC.Unit.State (lookupUnitId)
+
+import GHC.Utils.Exception
+import GHC.Utils.Error
+import GHC.Utils.Logger
+
+import System.Directory
+import System.FilePath
+import System.IO
+import System.IO.Error  ( isEOFError )
+import Control.Monad    ( when, unless )
+import Data.Foldable (traverse_)
+import Data.IORef
+import qualified Data.Set as Set
+import GHC.Driver.Make (downsweep)
+import GHC.Types.Error (unionManyMessages)
+import GHC.Unit (homeUnitId)
+
+#if __GLASGOW_HASKELL__ < 910
+import GHC.Utils.Panic.Plain
+#endif
+
+#if !defined(MWB)
+depJSON :: DynFlags -> Maybe FilePath
+depJSON _ = Nothing
+
+ms_opts :: ModSummary -> [String]
+ms_opts _ = []
+
+#endif
+
+-----------------------------------------------------------------
+--
+--              The main function
+--
+-----------------------------------------------------------------
+
+doMkDependHS :: GhcMonad m => [FilePath] -> m ModuleGraph
+doMkDependHS srcs = do
+    logger <- getLogger
+
+    -- Initialisation
+    dflags0 <- GHC.getSessionDynFlags
+
+    -- We kludge things a bit for dependency generation. Rather than
+    -- generating dependencies for each way separately, we generate
+    -- them once and then duplicate them for each way's osuf/hisuf.
+    -- We therefore do the initial dependency generation with an empty
+    -- way and .o/.hi extensions, regardless of any flags that might
+    -- be specified.
+    let dflags1 = dflags0
+            { targetWays_ = Set.empty
+            , hiSuf_      = "hi"
+            , objectSuf_  = "o"
+            }
+    -- GHC.setSessionDynFlags dflags1
+
+    -- If no suffix is provided, use the default -- the empty one
+    let dflags = if null (depSuffixes dflags0)
+                 then dflags0 { depSuffixes = [""] }
+                 else dflags0
+
+    tmpfs <- hsc_tmpfs <$> getSession
+    files <- liftIO $ beginMkDependHS logger tmpfs dflags
+
+    -- Do the downsweep to find all the modules
+    targets <- mapM (\s -> GHC.guessTarget s Nothing Nothing) srcs
+    GHC.setTargets targets
+    let excl_mods = depExcludeMods dflags
+    (errs, graph_nodes) <- withSession \ hsc_env -> liftIO $ downsweep hsc_env [] excl_mods True
+    let msgs = unionManyMessages errs
+    unless (isEmptyMessages msgs) $ throwErrors (fmap GhcDriverMessage msgs)
+    let module_graph = mkModuleGraph graph_nodes
+    -- Sort into dependency order
+    -- There should be no cycles
+    let sorted = GHC.topSortModuleGraph False module_graph Nothing
+
+    -- Print out the dependencies if wanted
+    liftIO $ debugTraceMsg logger 2 (text "Module dependencies" $$ ppr sorted)
+
+    -- Process them one by one, dumping results into makefile
+    -- and complaining about cycles
+    hsc_env <- getSession
+    root <- liftIO getCurrentDirectory
+    let excl_mods = depExcludeMods dflags
+    mapM_ (liftIO . processDeps dflags hsc_env excl_mods root (mkd_tmp_hdl files) (mkd_dep_json files)) sorted
+
+    -- If -ddump-mod-cycles, show cycles in the module graph
+    liftIO $ dumpModCycles logger module_graph
+
+    -- Tidy up
+    liftIO $ endMkDependHS logger files
+    pure module_graph
+
+-----------------------------------------------------------------
+--
+--              beginMkDependHs
+--      Create a temporary file,
+--      find the Makefile,
+--      slurp through it, etc
+--
+-----------------------------------------------------------------
+
+data MkDepFiles
+  = MkDep { mkd_make_file :: FilePath,          -- Name of the makefile
+            mkd_make_hdl  :: Maybe Handle,      -- Handle for the open makefile
+             -- | Output interface for the -dep-json file
+            mkd_dep_json  :: !(Maybe (JsonOutput DepJSON)),
+            mkd_tmp_file  :: FilePath,          -- Name of the temporary file
+            mkd_tmp_hdl   :: Handle }           -- Handle of the open temporary file
+
+beginMkDependHS :: Logger -> TmpFs -> DynFlags -> IO MkDepFiles
+beginMkDependHS logger tmpfs dflags = do
+        -- open a new temp file in which to stuff the dependency info
+        -- as we go along.
+  tmp_file <- newTempName logger tmpfs (tmpDir dflags) TFL_CurrentModule "dep"
+  tmp_hdl <- openFile tmp_file WriteMode
+
+        -- open the makefile
+  let makefile = depMakefile dflags
+  exists <- doesFileExist makefile
+  mb_make_hdl <-
+        if not exists
+        then return Nothing
+        else do
+           makefile_hdl <- openFile makefile ReadMode
+
+                -- slurp through until we get the magic start string,
+                -- copying the contents into dep_makefile
+           let slurp = do
+                l <- hGetLine makefile_hdl
+                if (l == depStartMarker)
+                        then return ()
+                        else do hPutStrLn tmp_hdl l; slurp
+
+                -- slurp through until we get the magic end marker,
+                -- throwing away the contents
+           let chuck = do
+                l <- hGetLine makefile_hdl
+                if (l == depEndMarker)
+                        then return ()
+                        else chuck
+
+           catchIO slurp
+                (\e -> if isEOFError e then return () else ioError e)
+           catchIO chuck
+                (\e -> if isEOFError e then return () else ioError e)
+
+           return (Just makefile_hdl)
+
+  dep_json_ref <- mkJsonOutput initDepJSON (depJSON dflags)
+
+        -- write the magic marker into the tmp file
+  hPutStrLn tmp_hdl depStartMarker
+
+  return (MkDep { mkd_make_file = makefile, mkd_make_hdl = mb_make_hdl,
+                  mkd_dep_json = dep_json_ref,
+                  mkd_tmp_file  = tmp_file, mkd_tmp_hdl  = tmp_hdl})
+
+-----------------------------------------------------------------
+--
+--              processDeps
+--
+-----------------------------------------------------------------
+
+processDeps :: DynFlags
+            -> HscEnv
+            -> [ModuleName]
+            -> FilePath
+            -> Handle           -- Write dependencies to here
+            -> Maybe (JsonOutput DepJSON)
+            -> SCC ModuleGraphNode
+            -> IO ()
+-- Write suitable dependencies to handle
+-- Always:
+--                      this.o : this.hs
+--
+-- If the dependency is on something other than a .hi file:
+--                      this.o this.p_o ... : dep
+-- otherwise
+--                      this.o ...   : dep.hi
+--                      this.p_o ... : dep.p_hi
+--                      ...
+-- (where .o is $osuf, and the other suffixes come from
+-- the cmdline -s options).
+--
+-- For {-# SOURCE #-} imports the "hi" will be "hi-boot".
+
+processDeps dflags _ _ _ _ _ (CyclicSCC nodes)
+  =     -- There shouldn't be any cycles; report them
+    throwGhcExceptionIO $ ProgramError $
+      showSDoc dflags $ GHC.cyclicModuleErr nodes
+
+processDeps dflags _ _ _ _ _ (AcyclicSCC (InstantiationNode _uid node))
+  =     -- There shouldn't be any backpack instantiations; report them as well
+    throwGhcExceptionIO $ ProgramError $
+      showSDoc dflags $
+        vcat [ text "Unexpected backpack instantiation in dependency graph while constructing Makefile:"
+             , nest 2 $ ppr node ]
+processDeps _dflags_ _ _ _ _ _ (AcyclicSCC (LinkNode {})) = return ()
+
+processDeps dflags hsc_env excl_mods root hdl m_dep_json (AcyclicSCC (ModuleNode _ node))
+  | moduleUnitId (ms_mod node) /= homeUnitId (hsc_home_unit hsc_env)
+  = pure ()
+  | otherwise = do
+  pp <- preprocessor
+  deps <- fmap concat $ sequence $
+    [cpp_deps | depIncludeCppDeps dflags] ++ [
+      import_deps IsBoot (ms_srcimps node),
+      import_deps NotBoot (ms_imps node)
+    ]
+  updateJson m_dep_json (updateDepJSON include_pkg_deps pp dep_node deps)
+  writeDependencies include_pkg_deps root hdl extra_suffixes dep_node deps
+  where
+    extra_suffixes = depSuffixes dflags
+    include_pkg_deps = depIncludePkgDeps dflags
+    src_file = msHsFilePath node
+    dep_node =
+      DepNode {
+        dn_mod = ms_mod node,
+        dn_src = src_file,
+        dn_obj = msObjFilePath node,
+        dn_hi = msHiFilePath node,
+        dn_boot = isBootSummary node,
+        dn_options = Set.fromList (ms_opts node)
+      }
+
+    preprocessor
+      | Just src <- ml_hs_file (ms_location node)
+      = runPipeline (hsc_hooks hsc_env) $ do
+        let (_, suffix) = splitExtension src
+            lit | Unlit _ <- startPhase suffix = True
+                | otherwise = False
+            pipe_env = mkPipeEnv StopPreprocess src Nothing NoOutputFile
+        unlit_fn <- if lit then use (T_Unlit pipe_env hsc_env src) else pure src
+        (dflags1, _, _) <- use (T_FileArgs hsc_env unlit_fn)
+        let pp = pgm_F dflags1
+        pure (if null pp then global_preprocessor else Just pp)
+      | otherwise
+      = pure global_preprocessor
+
+    global_preprocessor
+      | let pp = pgm_F dflags
+      , not (null pp)
+      = Just pp
+      | otherwise
+      = Nothing
+
+    -- Emit a dependency for each CPP import
+    -- CPP deps are discovered in the module parsing phase by parsing
+    -- comment lines left by the preprocessor.
+    -- Note that GHC.parseModule may throw an exception if the module
+    -- fails to parse, which may not be desirable (see #16616).
+    cpp_deps = do
+      session <- Session <$> newIORef hsc_env
+      parsedMod <- reflectGhc (GHC.parseModule node) session
+      pure (DepCpp <$> GHC.pm_extra_src_files parsedMod)
+
+    -- Emit a dependency for each import
+    import_deps is_boot idecls =
+      sequence [
+        findDependency hsc_env loc mb_pkg mod is_boot
+        | (mb_pkg, L loc mod) <- idecls
+        , mod `notElem` excl_mods
+        ]
+
+
+findDependency  :: HscEnv
+                -> SrcSpan
+                -> PkgQual              -- package qualifier, if any
+                -> ModuleName           -- Imported module
+                -> IsBootInterface      -- Source import
+                -> IO Dep
+findDependency hsc_env srcloc pkg imp dep_boot = do
+  -- Find the module; this will be fast because
+  -- we've done it once during downsweep
+  findImportedModule hsc_env imp pkg >>= \case
+    Found loc dep_mod ->
+      pure DepHi {
+        dep_mod,
+        dep_path = ml_hi_file loc,
+        dep_unit = lookupUnitId (hsc_units hsc_env) (moduleUnitId dep_mod),
+        dep_local,
+        dep_boot
+      }
+      where
+        dep_local = isJust (ml_hs_file loc) && moduleUnitId dep_mod == homeUnitId (hsc_home_unit hsc_env)
+
+    fail ->
+      throwOneError $
+      mkPlainErrorMsgEnvelope srcloc $
+      GhcDriverMessage $
+      DriverInterfaceError $
+      Can'tFindInterface (cannotFindModule hsc_env imp fail) $
+      LookingForModule imp dep_boot
+
+writeDependencies ::
+  Bool ->
+  FilePath ->
+  Handle ->
+  [FilePath] ->
+  DepNode ->
+  [Dep] ->
+  IO ()
+writeDependencies include_pkgs root hdl suffixes node deps =
+  traverse_ write tasks
+  where
+    tasks = source_dep : boot_dep ++ concatMap import_dep deps
+
+    -- Emit std dependency of the object(s) on the source file
+    -- Something like       A.o : A.hs
+    source_dep = (obj_files, dn_src)
+
+    -- add dependency between objects and their corresponding .hi-boot
+    -- files if the module has a corresponding .hs-boot file (#14482)
+    boot_dep
+      | IsBoot <- dn_boot
+      = [([obj], hi) | (obj, hi) <- zip (suffixed (removeBootSuffix dn_obj)) (suffixed dn_hi)]
+      | otherwise
+      = []
+
+    -- Add one dependency for each suffix;
+    -- e.g.         A.o   : B.hi
+    --              A.x_o : B.x_hi
+    import_dep = \case
+      DepHi {dep_path, dep_boot, dep_unit}
+        | isNothing dep_unit || include_pkgs
+        , let path = addBootSuffix_maybe dep_boot dep_path
+        -> [([obj], hi) | (obj, hi) <- zip obj_files (suffixed path)]
+
+        | otherwise
+        -> []
+
+      DepCpp {dep_path} -> [(obj_files, dep_path)]
+
+    write (from, to) = writeDependency root hdl from to
+
+    obj_files = suffixed dn_obj
+
+    suffixed f = insertSuffixes f suffixes
+
+    DepNode {dn_src, dn_obj, dn_hi, dn_boot} = node
+
+-----------------------------
+writeDependency :: FilePath -> Handle -> [FilePath] -> FilePath -> IO ()
+-- (writeDependency r h [t1,t2] dep) writes to handle h the dependency
+--      t1 t2 : dep
+writeDependency root hdl targets dep
+  = do let -- We need to avoid making deps on
+           --     c:/foo/...
+           -- on cygwin as make gets confused by the :
+           -- Making relative deps avoids some instances of this.
+           dep' = makeRelative root dep
+           forOutput = escapeSpaces . reslash Forwards . normalise
+           output = unwords (map forOutput targets) ++ " : " ++ forOutput dep'
+       hPutStrLn hdl output
+
+-----------------------------
+insertSuffixes
+        :: FilePath     -- Original filename;   e.g. "foo.o"
+        -> [String]     -- Suffix prefixes      e.g. ["x_", "y_"]
+        -> [FilePath]   -- Zapped filenames     e.g. ["foo.x_o", "foo.y_o"]
+        -- Note that the extra bit gets inserted *before* the old suffix
+        -- We assume the old suffix contains no dots, so we know where to
+        -- split it
+insertSuffixes file_name extras
+  = [ basename <.> (extra ++ suffix) | extra <- extras ]
+  where
+    (basename, suffix) = case splitExtension file_name of
+                         -- Drop the "." from the extension
+                         (b, s) -> (b, drop 1 s)
+
+
+-----------------------------------------------------------------
+--
+--              endMkDependHs
+--      Complete the makefile, close the tmp file etc
+--
+-----------------------------------------------------------------
+
+endMkDependHS :: Logger -> MkDepFiles -> IO ()
+
+endMkDependHS logger
+   (MkDep { mkd_make_file = makefile, mkd_make_hdl = makefile_hdl,
+            mkd_dep_json,
+            mkd_tmp_file = tmp_file, mkd_tmp_hdl = tmp_hdl })
+  = do
+  -- write the magic marker into the tmp file
+  hPutStrLn tmp_hdl depEndMarker
+
+  case makefile_hdl of
+     Nothing  -> return ()
+     Just hdl -> do
+        -- slurp the rest of the original makefile and copy it into the output
+        SysTools.copyHandle hdl tmp_hdl
+        hClose hdl
+
+  hClose tmp_hdl  -- make sure it's flushed
+
+        -- Create a backup of the original makefile
+  when (isJust makefile_hdl) $ do
+    showPass logger ("Backing up " ++ makefile)
+    SysTools.copyFile makefile (makefile++".bak")
+
+        -- Copy the new makefile in place
+  showPass logger "Installing new makefile"
+  SysTools.copyFile tmp_file makefile
+
+  -- Write the dependency and option data to a json file if the corresponding
+  -- flags were specified.
+  writeJsonOutput mkd_dep_json
+
+
+-----------------------------------------------------------------
+--              Module cycles
+-----------------------------------------------------------------
+
+dumpModCycles :: Logger -> ModuleGraph -> IO ()
+dumpModCycles logger module_graph
+  | not (logHasDumpFlag logger Opt_D_dump_mod_cycles)
+  = return ()
+
+  | null cycles
+  = putMsg logger (text "No module cycles")
+
+  | otherwise
+  = putMsg logger (hang (text "Module cycles found:") 2 pp_cycles)
+  where
+    topoSort = GHC.topSortModuleGraph True module_graph Nothing
+
+    cycles :: [[ModuleGraphNode]]
+    cycles =
+      [ c | CyclicSCC c <- topoSort ]
+
+    pp_cycles = vcat [ (text "---------- Cycle" <+> int n <+> text "----------")
+                        $$ pprCycle c $$ blankLine
+                     | (n,c) <- [1..] `zip` cycles ]
+
+pprCycle :: [ModuleGraphNode] -> SDoc
+-- Print a cycle, but show only the imports within the cycle
+pprCycle summaries = pp_group (CyclicSCC summaries)
+  where
+    cycle_mods :: [ModuleName]  -- The modules in this cycle
+    cycle_mods = map (moduleName . ms_mod) [ms | ModuleNode _ ms <- summaries]
+
+    pp_group :: SCC ModuleGraphNode -> SDoc
+    pp_group (AcyclicSCC (ModuleNode _ ms)) = pp_ms ms
+    pp_group (AcyclicSCC _) = empty
+    pp_group (CyclicSCC mss)
+        = assert (not (null boot_only)) $
+                -- The boot-only list must be non-empty, else there would
+                -- be an infinite chain of non-boot imports, and we've
+                -- already checked for that in processModDeps
+          pp_ms loop_breaker $$ vcat (map pp_group groups)
+        where
+          (boot_only, others) = partition is_boot_only mss
+          is_boot_only (ModuleNode _ ms) = not (any in_group (map snd (ms_imps ms)))
+          is_boot_only  _ = False
+          in_group (L _ m) = m `elem` group_mods
+          group_mods = map (moduleName . ms_mod) [ms | ModuleNode _ ms <- mss]
+
+          loop_breaker = head ([ms | ModuleNode _ ms  <- boot_only])
+          all_others   = tail boot_only ++ others
+          groups =
+            GHC.topSortModuleGraph True (mkModuleGraph all_others) Nothing
+
+    pp_ms summary = text mod_str <> text (take (20 - length mod_str) (repeat ' '))
+                       <+> (pp_imps empty (map snd (ms_imps summary)) $$
+                            pp_imps (text "{-# SOURCE #-}") (map snd (ms_srcimps summary)))
+        where
+          mod_str = moduleNameString (moduleName (ms_mod summary))
+
+    pp_imps :: SDoc -> [Located ModuleName] -> SDoc
+    pp_imps _    [] = empty
+    pp_imps what lms
+        = case [m | L _ m <- lms, m `elem` cycle_mods] of
+            [] -> empty
+            ms -> what <+> text "imports" <+>
+                                pprWithCommas ppr ms
+
+-----------------------------------------------------------------
+--
+--              Flags
+--
+-----------------------------------------------------------------
+
+depStartMarker, depEndMarker :: String
+depStartMarker = "# DO NOT DELETE: Beginning of Haskell dependencies"
+depEndMarker   = "# DO NOT DELETE: End of Haskell dependencies"
diff --git a/plugin/src/Internal/MakeFile/JSON.hs b/plugin/src/Internal/MakeFile/JSON.hs
new file mode 100644
index 00000000..15370304
--- /dev/null
+++ b/plugin/src/Internal/MakeFile/JSON.hs
@@ -0,0 +1,240 @@
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingVia #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE NoFieldSelectors #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+{-# LANGUAGE CPP #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
+module Internal.MakeFile.JSON
+  ( writeJSONFile,
+    JsonOutput (..),
+    mkJsonOutput,
+    updateJson,
+    writeJsonOutput,
+    DepJSON,
+    DepNode (..),
+    Dep (..),
+    initDepJSON,
+    updateDepJSON,
+  )
+where
+
+import Data.Foldable (traverse_)
+import Data.IORef
+import qualified Data.Map.Strict as Map
+import qualified Data.Semigroup as Semigroup
+import qualified Data.Set as Set
+import GHC.Data.FastString (unpackFS)
+import GHC.Generics (Generic, Generically (Generically))
+import GHC.Prelude
+import GHC.Unit
+import GHC.Utils.Json
+import GHC.Utils.Misc
+import GHC.Utils.Outputable
+import System.FilePath (normalise)
+
+#if !defined(MWB)
+import GHC.Data.FastString (lexicalCompareFS)
+
+instance Ord PackageId where
+    PackageId p1 `compare` PackageId p2 = p1 `lexicalCompareFS` p2
+
+#endif
+
+--------------------------------------------------------------------------------
+-- Output helpers
+--------------------------------------------------------------------------------
+
+writeJSONFile :: ToJson a => a -> FilePath -> IO ()
+writeJSONFile doc p = do
+  withAtomicRename p
+    $ \tmp -> writeFile tmp $ showSDocUnsafe $ renderJSON $ json doc
+
+--------------------------------------------------------------------------------
+-- Output interface for json dumps
+--------------------------------------------------------------------------------
+
+-- | Resources for a json dump option, used in "GHC.Driver.MakeFile".
+-- The flag @-dep-json@ add an additional output target for dependency
+-- diagnostics.
+data JsonOutput a =
+  JsonOutput {
+    -- | This ref is updated in @processDeps@ incrementally, using a
+    -- flag-specific type.
+    json_ref :: IORef a,
+
+    -- | The output file path specified as argument to the flag.
+    json_path :: FilePath
+  }
+
+-- | Allocate an 'IORef' with the given function if the 'FilePath' is 'Just',
+-- indicating that the userspecified @-*-json@.
+mkJsonOutput ::
+  IO (IORef a) ->
+  Maybe FilePath ->
+  IO (Maybe (JsonOutput a))
+mkJsonOutput mk_ref =
+  traverse $ \ json_path -> do
+    json_ref <- mk_ref
+    pure JsonOutput {json_ref, json_path}
+
+-- | Update the dump data in 'json_ref' if the output target is present.
+updateJson :: Maybe (JsonOutput a) -> (a -> a) -> IO ()
+updateJson out f = traverse_ (\ JsonOutput {json_ref} -> modifyIORef' json_ref f) out
+
+-- | Write a json object to the flag-dependent file if the output target is
+-- present.
+writeJsonOutput ::
+  ToJson a =>
+  Maybe (JsonOutput a) ->
+  IO ()
+writeJsonOutput =
+  traverse_ $ \ JsonOutput {json_ref, json_path} -> do
+    payload <- readIORef json_ref
+    writeJSONFile payload json_path
+
+--------------------------------------------------------------------------------
+-- Types abstracting over json and Makefile
+--------------------------------------------------------------------------------
+
+data DepNode =
+  DepNode {
+    dn_mod :: Module,
+    dn_src :: FilePath,
+    dn_obj :: FilePath,
+    dn_hi :: FilePath,
+    dn_boot :: IsBootInterface,
+    dn_options :: Set.Set String
+  }
+
+data Dep =
+  DepHi {
+    dep_mod :: Module,
+    dep_path :: FilePath,
+    dep_unit :: Maybe UnitInfo,
+    dep_local :: Bool,
+    dep_boot :: IsBootInterface
+  }
+  |
+  DepCpp {
+    dep_path :: FilePath
+  }
+
+--------------------------------------------------------------------------------
+-- Payload for -dep-json
+--------------------------------------------------------------------------------
+
+newtype PackageDeps =
+  PackageDeps (Map.Map (String, UnitId, PackageId) (Set.Set ModuleName))
+  deriving newtype (Monoid)
+
+instance Semigroup PackageDeps where
+  PackageDeps l <> PackageDeps r = PackageDeps (Map.unionWith (Semigroup.<>) l r)
+
+data Deps =
+  Deps {
+    sources :: Set.Set FilePath,
+    modules :: (Set.Set ModuleName, Set.Set ModuleName),
+    packages :: PackageDeps,
+    cpp :: Set.Set FilePath,
+    options :: Set.Set String,
+    preprocessor :: Maybe FilePath
+  }
+  deriving stock (Generic)
+  deriving (Semigroup, Monoid) via (Generically Deps)
+
+data ModuleDeps =
+  ModuleDeps {
+    source :: Deps,
+    boot :: Maybe Deps
+  }
+  deriving stock (Generic)
+  deriving (Semigroup, Monoid) via (Generically ModuleDeps)
+
+newtype DepJSON = DepJSON (Map.Map ModuleName ModuleDeps)
+
+instance ToJson DepJSON where
+  json (DepJSON m) =
+    JSObject [
+      (moduleNameString target, module_deps md)
+      | (target, md) <- Map.toList m
+    ]
+    where
+      module_deps ModuleDeps {source, boot} =
+        JSObject (("boot", maybe JSNull (JSObject . deps) boot) : deps source)
+
+      deps Deps {packages = PackageDeps packages, ..} =
+        [
+          ("sources", array sources normalise),
+          ("modules", array (fst modules) moduleNameString),
+          ("modules-boot", array (snd modules) moduleNameString),
+          ("packages",
+            JSArray [
+              package name unit_id package_id mods |
+              ((name, unit_id, package_id), mods) <- Map.toList packages
+            ]
+          ),
+          ("cpp", array cpp id),
+          ("options", array options id),
+          ("preprocessor", maybe JSNull JSString preprocessor)
+        ]
+
+      package name unit_id (PackageId package_id) mods =
+        JSObject [
+          ("id", JSString (unitIdString unit_id)),
+          ("name", JSString name),
+          ("package-id", JSString (unpackFS package_id)),
+          ("modules", array mods moduleNameString)
+        ]
+
+      array values render = JSArray (fmap (JSString . render) (Set.toList values))
+
+initDepJSON :: IO (IORef DepJSON)
+initDepJSON = newIORef $ DepJSON Map.empty
+
+insertDepJSON :: ModuleName -> IsBootInterface -> Deps -> DepJSON -> DepJSON
+insertDepJSON target is_boot dep (DepJSON m0) =
+  DepJSON $ Map.insertWith (Semigroup.<>) target new m0
+  where
+    new
+      | IsBoot <- is_boot = mempty {boot = Just dep}
+      | otherwise = mempty {source = dep}
+
+updateDepJSON :: Bool -> Maybe FilePath -> DepNode -> [Dep] -> DepJSON -> DepJSON
+updateDepJSON include_pkgs preprocessor DepNode {..} deps =
+  insertDepJSON (moduleName dn_mod) dn_boot payload
+  where
+    payload = node_data Semigroup.<> foldMap dep deps
+
+    node_data =
+      mempty {
+        sources = Set.singleton dn_src,
+        preprocessor,
+        options = dn_options
+      }
+
+    dep = \case
+      DepHi {dep_mod, dep_local, dep_unit, dep_boot}
+        | dep_local
+        , let set = Set.singleton (moduleName dep_mod)
+              value | IsBoot <- dep_boot = (Set.empty, set)
+                    | otherwise = (set, Set.empty)
+        -> mempty {modules = value}
+
+        | include_pkgs
+        , Just unit <- dep_unit
+        , let PackageName nameFS = unitPackageName unit
+              name = unpackFS nameFS
+              withLibName (PackageName c) = name ++ ":" ++ unpackFS c
+              lname = maybe name withLibName (unitComponentName unit)
+              key = (lname, unitId unit, unitPackageId unit)
+        -> mempty {packages = PackageDeps (Map.singleton key (Set.singleton (moduleName dep_mod)))}
+
+        | otherwise
+        -> mempty
+
+      DepCpp {dep_path} ->
+        mempty {cpp = Set.singleton dep_path}
diff --git a/plugin/src/Internal/Metadata.hs b/plugin/src/Internal/Metadata.hs
new file mode 100644
index 00000000..0bde45ce
--- /dev/null
+++ b/plugin/src/Internal/Metadata.hs
@@ -0,0 +1,53 @@
+module Internal.Metadata where
+
+import Control.Concurrent (readMVar)
+import Control.Monad.IO.Class (liftIO)
+import Data.Maybe (fromMaybe)
+import GHC (DynFlags (..), GhcMode (..))
+import GHC.Driver.Env (HscEnv (..), hscUpdateFlags)
+import GHC.Driver.Monad (modifySession, modifySessionM, withTempSession)
+import GHC.Platform.Ways (Way (WayDyn), addWay)
+import GHC.Runtime.Loader (initializeSessionPlugins)
+import GHC.Unit.Env (UnitEnv (..), unitEnv_union)
+import Internal.Cache (Cache (..), Target (..), mergeHugs, newFinderCache, updateModuleGraph)
+import Internal.MakeFile (doMkDependHS)
+import Internal.Session (Env (..), runSession, withGhcInSession)
+
+-- | Copy the cached unit env and module graph to the given session.
+restoreEnv :: Env -> HscEnv -> IO HscEnv
+restoreEnv env hsc_env = do
+  cache <- readMVar env.cache
+  pure $ maybe id restoreMg cache.moduleGraph $ maybe hsc_env restore cache.hug
+  where
+    restoreMg new e = e {hsc_mod_graph = new}
+
+    restore hug =
+      hsc_env {hsc_unit_env = hsc_env.hsc_unit_env {ue_home_unit_graph = unitEnv_union mergeHugs hug current}}
+
+    current = hsc_env.hsc_unit_env.ue_home_unit_graph
+
+-- | Run downsweep and merge the resulting module graph into the cached graph.
+-- This is executed for the metadata step, which natively only calls 'doMkDependHS'.
+-- Since that function doesn't give us access to the module graph in its original shape, we inline it into this project
+-- to exfiltrate the graph.
+-- This has @WayDyn@ hardcoded for now, but it should be adapted to Buck's build configuration.
+-- This is usually not necessary (in fact, 'doMkDependHS' clears the target ways) but since we're keeping the module
+-- graph the target way will be reflected in the stored @ModSummary@ nodes.
+--
+-- Before downsweep, we also create a fresh @Finder@ for some reason and restore the previous unit env so dependencies
+-- are visible.
+computeMetadata :: Env -> IO Bool
+computeMetadata env =
+  fmap (fromMaybe False) $ runSession False env $ withGhcInSession env \ srcs -> do
+    initializeSessionPlugins
+    cache <- liftIO $ readMVar env.cache
+    module_graph <- withTempSession (hscUpdateFlags (\ d -> d { targetWays_ = addWay WayDyn (targetWays_ d) })) do
+      modifySessionM (liftIO . restoreEnv env)
+      modifySessionM \ hsc_env -> do
+        hsc_FC <- liftIO $ newFinderCache env.cache cache (Target "metadata")
+        pure hsc_env {hsc_FC}
+      modifySession $ hscUpdateFlags \ d -> d {ghcMode = MkDepend}
+      module_graph <- doMkDependHS (fst <$> srcs)
+      pure module_graph
+    liftIO $ updateModuleGraph env.cache module_graph
+    pure (Just True)

From 614ff1d2f9637bf1037651092e29e36b8e4f015d Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Wed, 9 Apr 2025 01:05:32 -0700
Subject: [PATCH 13/16] Debug show functions added.

ModGraph, Eps, Hpt, HUG are now printable for debug purpose.

Change-Id: I2f4a609708bb4fcfff38ff897e6979ea03dd3ea4
---
 buck-worker/Main.hs                       |  27 +++--
 plugin/ghc-persistent-worker-plugin.cabal |   1 +
 plugin/src/BUCK                           |   1 +
 plugin/src/Internal/Cache.hs              |  24 ++++-
 plugin/src/Internal/Debug.hs              | 123 ++++++++++++++++++++++
 plugin/src/Internal/Session.hs            |   2 -
 6 files changed, 158 insertions(+), 20 deletions(-)
 create mode 100644 plugin/src/Internal/Debug.hs

diff --git a/buck-worker/Main.hs b/buck-worker/Main.hs
index 60dfe567..603d559a 100644
--- a/buck-worker/Main.hs
+++ b/buck-worker/Main.hs
@@ -4,20 +4,19 @@ module Main where
 
 import BuckArgs (BuckArgs (abiOut), CompileResult (..), Mode (..), parseBuckArgs, toGhcArgs, writeResult)
 import qualified BuckArgs (BuckArgs (mode))
-import BuckWorker (ExecuteCommand, ExecuteCommand'EnvironmentEntry, ExecuteEvent, ExecuteResponse, Worker (..), Instrument (..))
+import BuckWorker (
+  ExecuteCommand,
+  ExecuteCommand'EnvironmentEntry,
+  ExecuteEvent,
+  ExecuteResponse,
+  Instrument (..),
+  Worker (..),
+  )
 import Control.Concurrent (MVar, modifyMVar_, newMVar)
 import Control.Concurrent.Async (async, cancel, wait)
-import Control.Concurrent.Chan (Chan, dupChan, readChan, newChan, writeChan)
-import Control.Exception (
-  Exception (displayException),
-  SomeException (SomeException),
-  bracket,
-  finally,
-  onException,
-  throwIO,
-  try,
-  )
-import Control.Monad (foldM, void, when, forever)
+import Control.Concurrent.Chan (Chan, dupChan, newChan, readChan, writeChan)
+import Control.Exception (Exception (..), SomeException (..), bracket, finally, onException, throwIO, try)
+import Control.Monad (foldM, forever, void, when)
 import Control.Monad.IO.Class (liftIO)
 import Data.Foldable (for_)
 import Data.Functor ((<&>))
@@ -31,7 +30,6 @@ import GHC (DynFlags (..), Ghc, getSession)
 import GHC.Driver.DynFlags (GhcMode (..))
 import GHC.Driver.Env (hscUpdateFlags)
 import GHC.Driver.Monad (modifySession)
-import GHC (Ghc, getSession)
 import GHC.IO.Handle.FD (withFileBlocking)
 import GHC.Stats (GCDetails (..), RTSStats (..), getRTSStats)
 import Internal.AbiHash (AbiHash (..), showAbiHash)
@@ -41,7 +39,7 @@ import Internal.CompileHpt (compileModuleWithDepsInHpt)
 import Internal.Log (dbg, logFlush, newLog)
 import Internal.Metadata (computeMetadata)
 import Internal.Session (Env (..), withGhcMhu)
-import Network.GRPC.Client (Connection, Server (ServerUnix), recvNextOutput, sendFinalInput, withConnection, withRPC)
+import Network.GRPC.Client (Connection, Server (..), recvNextOutput, sendFinalInput, withConnection, withRPC)
 import Network.GRPC.Common (NextElem (..), Proxy (..), def)
 import Network.GRPC.Common.Protobuf (Proto, Protobuf, defMessage, (%~), (&), (.~), (^.))
 import Network.GRPC.Server.Protobuf (ProtobufMethodsOf)
@@ -311,7 +309,6 @@ primarySocketDiscoveryIn dir = PrimarySocketDiscoveryPath (dir </> "primary")
 runLocalGhc :: WorkerMode -> ServerSocketPath -> Maybe InstrumentSocketPath -> IO ()
 runLocalGhc mode socket minstr = do
   dbg ("Starting ghc server on " ++ socket.path)
-
   cache <-
     case mode of
       WorkerMakeMode ->
diff --git a/plugin/ghc-persistent-worker-plugin.cabal b/plugin/ghc-persistent-worker-plugin.cabal
index 7fe19c77..48031957 100644
--- a/plugin/ghc-persistent-worker-plugin.cabal
+++ b/plugin/ghc-persistent-worker-plugin.cabal
@@ -30,6 +30,7 @@ library
                       Internal.Cache
                       Internal.Compile
                       Internal.CompileHpt
+                      Internal.Debug
                       Internal.Error
                       Internal.Log
                       Internal.MakeFile
diff --git a/plugin/src/BUCK b/plugin/src/BUCK
index fc104550..c8abfc36 100644
--- a/plugin/src/BUCK
+++ b/plugin/src/BUCK
@@ -4,6 +4,7 @@
     "Cache",
     "Compile",
     "CompileHpt",
+    "Debug",
     "Error",
     "Log",
     "MakeFile",
diff --git a/plugin/src/Internal/Cache.hs b/plugin/src/Internal/Cache.hs
index 8a1a5f2d..a2bacbcc 100644
--- a/plugin/src/Internal/Cache.hs
+++ b/plugin/src/Internal/Cache.hs
@@ -8,6 +8,7 @@ import Control.Monad.IO.Class (MonadIO, liftIO)
 import Data.Bifunctor (first)
 import Data.Coerce (coerce)
 import Data.Foldable (for_)
+import Data.IORef (readIORef)
 import Data.List (sortBy)
 import qualified Data.Map.Strict as Map
 import Data.Map.Strict (Map, (!?))
@@ -30,7 +31,10 @@ import GHC.Types.Unique.FM (UniqFM, minusUFM, nonDetEltsUFM, sizeUFM)
 import GHC.Types.Unique.Supply (initUniqSupply)
 import GHC.Unit.Env (HomeUnitEnv (..), HomeUnitGraph, UnitEnv (..), unitEnv_union)
 import GHC.Unit.External (ExternalUnitCache (..), initExternalUnitCache)
-import GHC.Unit.Module.Env (emptyModuleEnv, moduleEnvKeys, plusModuleEnv)
+import GHC.Unit.Finder (InstalledFindResult (..))
+import GHC.Unit.Finder.Types (FinderCache (..))
+import GHC.Unit.Module.Env (InstalledModuleEnv, emptyModuleEnv, moduleEnvKeys, plusModuleEnv)
+import GHC.Unit.Module.Graph (ModuleGraph, unionMG)
 import qualified GHC.Utils.Outputable as Outputable
 import GHC.Utils.Outputable (SDoc, comma, doublePrec, fsep, hang, nest, punctuate, text, vcat, ($$), (<+>))
 import Internal.Log (Log, logd)
@@ -38,7 +42,7 @@ import System.Environment (lookupEnv)
 
 #if __GLASGOW_HASKELL__ >= 911 || defined(MWB)
 
-import Data.IORef (IORef, newIORef, readIORef)
+import Data.IORef (IORef, newIORef)
 import qualified Data.Map.Lazy as LazyMap
 import GHC.Fingerprint (Fingerprint, getFileHash)
 import GHC.IORef (atomicModifyIORef')
@@ -50,7 +54,13 @@ import GHC.Utils.Panic (panic)
 
 #else
 
-import GHC.Unit.Finder (FinderCache, initFinderCache)
+import GHC.Unit.Finder (initFinderCache)
+
+#endif
+
+#if defined(MWB)
+
+import GHC.Unit.Module.Graph (ModuleGraphNode (..), mgModSummaries', mkModuleGraph, mkNodeKey)
 
 #endif
 
@@ -199,6 +209,10 @@ emptyFinderState =
     files <- newIORef LazyMap.empty
     pure FinderState {modules, files}
 
+finderEnv :: FinderState -> IO (InstalledModuleEnv InstalledFindResult)
+finderEnv FinderState {modules} =
+  readIORef modules
+
 #else
 
 data FinderState =
@@ -212,6 +226,10 @@ emptyFinderState =
     cache <- initFinderCache
     pure FinderState {cache}
 
+finderEnv :: FinderState -> IO (InstalledModuleEnv InstalledFindResult)
+finderEnv FinderState {cache = FinderCache {fcModuleCache}} =
+  readIORef fcModuleCache
+
 #endif
 
 data CacheFeatures =
diff --git a/plugin/src/Internal/Debug.hs b/plugin/src/Internal/Debug.hs
new file mode 100644
index 00000000..98a8f333
--- /dev/null
+++ b/plugin/src/Internal/Debug.hs
@@ -0,0 +1,123 @@
+{-# language OverloadedStrings, CPP #-}
+
+module Internal.Debug where
+
+import qualified Data.Map.Strict as Map
+import GHC (DynFlags (..), mi_module)
+import GHC.Types.Unique.DFM (udfmToList)
+import GHC.Types.Unique.Map (nonDetEltsUniqMap)
+import GHC.Unit (UnitDatabase (..), UnitId, UnitState (..), homeUnitId, moduleEnvToList, unitPackageId)
+import GHC.Unit.Env (HomeUnitEnv (..), HomeUnitGraph, UnitEnv (..), UnitEnvGraph (..))
+import GHC.Unit.External (ExternalPackageState (..), eucEPS)
+import GHC.Unit.Home.ModInfo (HomeModInfo (..), HomePackageTable, hm_iface)
+import GHC.Unit.Module.Graph (ModuleGraph)
+import GHC.Utils.Outputable (Outputable, SDoc, comma, hang, hcat, ppr, punctuate, text, vcat, (<+>))
+import System.FilePath (takeDirectory, takeFileName)
+
+#if __GLASGOW_HASKELL__ < 911 && !defined(MWB)
+
+import Data.Foldable (toList)
+import GHC.Unit.Module.Graph (mgTransDeps)
+
+#else
+
+import GHC.Unit.Module.Graph (mgModSummaries')
+
+#endif
+
+entryD :: (SDoc, SDoc) -> SDoc
+entryD (k, v) = hang (hcat [k, ":"]) 2 v
+
+entry :: (String, SDoc) -> SDoc
+entry (k, v) = entryD (text k, v)
+
+entries :: [(String, SDoc)] -> SDoc
+entries = vcat . fmap entry
+
+showMap ::
+  Outputable a =>
+  (b -> SDoc) ->
+  [(a, b)] ->
+  SDoc
+showMap pprB m =
+  vcat [ppr from <+> text "->" <+> (pprB to) | (from, to) <- m]
+
+#if __GLASGOW_HASKELL__ < 911 && !defined(MWB)
+
+showModGraph :: ModuleGraph -> SDoc
+showModGraph g =
+  showMap (ppr . toList) (Map.toList (mgTransDeps g))
+
+#else
+
+showModGraph :: ModuleGraph -> SDoc
+showModGraph g =
+  ppr (mgModSummaries' g)
+
+#endif
+
+showEps :: ExternalPackageState -> IO SDoc
+showEps EPS {..} = do
+  pure $ entries $ [
+    ] ++ if False then [pit] else []
+  where
+    pit = ("pit", vcat [ppr m <+> ppr (mi_module iface) | (m, iface) <- moduleEnvToList eps_PIT])
+
+showUnitState :: UnitState -> SDoc
+showUnitState UnitState {..} =
+  entries $ [
+    ("homeUnitDepends", ppr homeUnitDepends)
+  ] ++
+  if False
+  then [("unitInfoMap", ppr (ppr . unitPackageId <$> nonDetEltsUniqMap unitInfoMap))]
+  else []
+
+showHomeUnitDflags :: DynFlags -> SDoc
+showHomeUnitDflags DynFlags {..} =
+  entries [
+    ("homeUnitId", ppr homeUnitId_)
+  ]
+
+showHpt :: HomePackageTable -> SDoc
+showHpt hpt =
+  hcat (punctuate comma [ppr (mi_module hm_iface) | (_, HomeModInfo {..}) <- udfmToList hpt])
+   -- <+> ppr hm_linkable
+
+showDbPath :: UnitDatabase UnitId -> SDoc
+showDbPath UnitDatabase {unitDatabasePath} =
+  text (takeFileName (takeDirectory unitDatabasePath))
+
+showHomeUnitEnvShort :: HomeUnitEnv -> SDoc
+showHomeUnitEnvShort HomeUnitEnv {..} =
+  entries [
+    ("deps", ppr homeUnitEnv_units.homeUnitDepends),
+    ("dbs", maybe (text "not loaded") (ppr . fmap showDbPath) homeUnitEnv_unit_dbs),
+    ("hpt", showHpt homeUnitEnv_hpt)
+  ]
+
+showHomeUnitEnv :: HomeUnitEnv -> SDoc
+showHomeUnitEnv HomeUnitEnv {..} =
+  entries [
+    ("units", showUnitState homeUnitEnv_units),
+    ("homeUnitEnv_unit_dbs", ppr homeUnitEnv_unit_dbs),
+    ("dflags", showHomeUnitDflags homeUnitEnv_dflags),
+    ("hpt", showHpt homeUnitEnv_hpt),
+    ("home_unit", ppr (homeUnitId <$> homeUnitEnv_home_unit))
+  ]
+
+showHugShort :: HomeUnitGraph -> SDoc
+showHugShort (UnitEnvGraph hug) =
+  vcat [entryD ((ppr k), (showHomeUnitEnvShort e)) | (k, e) <- Map.toList hug]
+
+showHug :: HomeUnitGraph -> SDoc
+showHug (UnitEnvGraph hug) =
+  vcat [entryD ((ppr k), (showHomeUnitEnv e)) | (k, e) <- Map.toList hug]
+
+showUnitEnv :: UnitEnv -> IO SDoc
+showUnitEnv UnitEnv {..} = do
+  eps <- showEps =<< eucEPS ue_eps
+  pure $ entries [
+    ("eps", eps),
+    ("hug", showHug ue_home_unit_graph),
+    ("current_unit", ppr ue_current_unit)
+    ]
diff --git a/plugin/src/Internal/Session.hs b/plugin/src/Internal/Session.hs
index 803c2d0c..8ff0f2b8 100644
--- a/plugin/src/Internal/Session.hs
+++ b/plugin/src/Internal/Session.hs
@@ -13,7 +13,6 @@ import GHC (
   DynFlags (..),
   Ghc,
   GhcLink (LinkBinary),
-  GhcMode (OneShot),
   Phase,
   getSessionDynFlags,
   parseDynamicFlags,
@@ -180,7 +179,6 @@ withGhcDefault :: Env -> (Target -> Ghc (Maybe (Maybe ModuleArtifacts, a))) -> I
 withGhcDefault env =
   withGhcUsingCache (withCache env.log env.args.workerTargetId env.cache) env
 
-
 -- | Command line args that have to be stored in the current home unit env.
 -- These are specified as a single program argument with their option argument, without whitespace in between.
 specificPrefixSwitches :: [String]

From 1d31a4e8a6a9c6ddd4e6c95ba0edb4e12d635b60 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Mon, 7 Apr 2025 15:39:39 -0700
Subject: [PATCH 14/16] Test code for compilation using HPT

Change-Id: Ie47a760fa34bf0402af12ec61863346ed0bec909
---
 buck-worker/buck-worker.cabal      |  18 ++
 buck-worker/lib/BuckArgs.hs        |   4 +-
 buck-worker/test/CompileHptTest.hs | 338 +++++++++++++++++++++++++++++
 buck-worker/test/Main.hs           |   6 +
 buck-worker/test/TestSetup.hs      | 175 +++++++++++++++
 5 files changed, 539 insertions(+), 2 deletions(-)
 create mode 100644 buck-worker/test/CompileHptTest.hs
 create mode 100644 buck-worker/test/Main.hs
 create mode 100644 buck-worker/test/TestSetup.hs

diff --git a/buck-worker/buck-worker.cabal b/buck-worker/buck-worker.cabal
index 28379e10..1bacb7aa 100644
--- a/buck-worker/buck-worker.cabal
+++ b/buck-worker/buck-worker.cabal
@@ -61,3 +61,21 @@ executable worker
     ghc-persistent-worker-plugin,
     text,
     buck-worker
+
+test-suite worker-test
+  import: all
+  type: exitcode-stdio-1.0
+  main-is: Main.hs
+  hs-source-dirs: test
+  other-modules:
+    CompileHptTest,
+    TestSetup
+  ghc-options: -threaded -rtsopts -with-rtsopts=-N -with-rtsopts=-T
+  build-depends:
+    directory,
+    filepath,
+    ghc,
+    ghc-persistent-worker-plugin,
+    temporary,
+    transformers,
+    typed-process
diff --git a/buck-worker/lib/BuckArgs.hs b/buck-worker/lib/BuckArgs.hs
index 0a95eec4..5ff0c0eb 100644
--- a/buck-worker/lib/BuckArgs.hs
+++ b/buck-worker/lib/BuckArgs.hs
@@ -104,8 +104,8 @@ options =
     withArg "--bin-exe" \ z a -> z {binPath = takeDirectory a : z.binPath},
     withArg "--worker-mode" \ z a -> z {mode = Just (parseMode a)},
     flag "--worker-multiplexer-custom" \ z -> z {multiplexerCustom = True},
-    ("-c", \ rest z -> Right (rest, z {mode = Just ModeCompile}))
-    -- skip "-c"
+    ("-c", \ rest z -> Right (rest, z {mode = Just ModeCompile})),
+    ("-M", \ rest z -> Right (rest, z {mode = Just ModeMetadata}))
   ]
   where
     addEnv z a = case z.envKey of
diff --git a/buck-worker/test/CompileHptTest.hs b/buck-worker/test/CompileHptTest.hs
new file mode 100644
index 00000000..7584d589
--- /dev/null
+++ b/buck-worker/test/CompileHptTest.hs
@@ -0,0 +1,338 @@
+module CompileHptTest where
+
+import Control.Concurrent (readMVar)
+import Control.Monad (when)
+import Control.Monad.IO.Class (liftIO)
+import Control.Monad.Trans.State.Strict (StateT, evalStateT, state)
+import Data.Char (toUpper)
+import Data.Foldable (for_, traverse_)
+import Data.Functor (void, (<&>))
+import Data.List (intercalate)
+import Data.Maybe (fromMaybe, isNothing)
+import qualified Data.Set as Set
+import Data.Set (Set)
+import GHC (
+  DynFlags (..),
+  Ghc,
+  GhcException (..),
+  GhcMode (..),
+  ModSummary (..),
+  ModuleGraph,
+  getSession,
+  guessTarget,
+  mgModSummaries,
+  mkModuleGraph,
+  ms_mod_name,
+  setTargets,
+  unLoc,
+  )
+import GHC.Driver.Env (HscEnv (..), hscUpdateFlags, hsc_home_unit)
+import GHC.Driver.Make (downsweep)
+import GHC.Driver.Monad (modifySession, modifySessionM, withTempSession)
+import GHC.Runtime.Loader (initializeSessionPlugins)
+import GHC.Types.Error (unionManyMessages)
+import GHC.Unit (homeUnitId, moduleUnitId)
+import GHC.Unit.Env (UnitEnv (..), unitEnv_union)
+import GHC.Unit.Finder (addHomeModuleToFinder)
+import GHC.Utils.Outputable (ppr, text, (<+>))
+import GHC.Utils.Panic (throwGhcExceptionIO)
+import Internal.Args (Args (..))
+import Internal.Cache (Cache (..), Target (..), mergeHugs, newFinderCache, updateModuleGraph)
+import Internal.CompileHpt (adaptHp, compileModuleWithDepsInHpt, initUnit)
+import Internal.Debug (showHugShort, showModGraph)
+import Internal.Log (dbg, dbgp, dbgs, newLog)
+import Internal.Session (Env (..), dummyLocation, withGhcMhu, withGhcInSession, withUnitSpecificOptions)
+import Prelude hiding (log)
+import System.Directory (listDirectory)
+import System.FilePath (dropExtension, takeDirectory, takeExtension, takeFileName, (</>))
+import TestSetup (Conf (..), UnitConf (..), UnitMod (..), withProject)
+
+debugState :: Bool
+debugState = False
+
+-- | Approximate synthetic reproduction of what happens when the metadata step is performed by the worker.
+loadModuleGraph :: Env -> UnitMod -> [String] -> Ghc ModuleGraph
+loadModuleGraph env UnitMod {src} specific = do
+  cache <- liftIO $ readMVar env.cache
+  (_, module_graph) <- withTempSession (maybe id restoreHug cache.hug . maybe id restoreModuleGraph cache.moduleGraph) do
+    modifySessionM \ hsc_env -> do
+      hsc_FC <- liftIO $ newFinderCache env.cache cache (Target "metadata")
+      pure hsc_env {hsc_FC}
+    dbg (unwords specific)
+    initUnit specific
+    when debugState do
+      mg <- hsc_mod_graph <$> getSession
+      dbgp ("existing graph:" <+> showModGraph mg)
+      dbg "hug:"
+      dbgp . showHugShort . ue_home_unit_graph . hsc_unit_env =<< getSession
+    names <- liftIO $ listDirectory dir
+    let srcs = [dir </> name | name <- names, takeExtension name == ".hs"]
+    targets <- mapM (\s -> guessTarget s Nothing Nothing) srcs
+    setTargets targets
+    hsc_env1 <- getSession
+    (errs, graph_nodes) <- liftIO $ downsweep hsc_env1 [] [] True
+    let
+      mod_graph = mkModuleGraph graph_nodes
+    pure (unionManyMessages errs, mod_graph)
+  when debugState do
+    dbgp (text "unit module graph:" <+> showModGraph module_graph)
+  liftIO $ updateModuleGraph env.cache module_graph
+  pure module_graph
+  where
+    dir = takeDirectory src
+
+    restoreModuleGraph mg e = e {hsc_mod_graph = mg}
+
+    restoreHug hug e =
+      e {hsc_unit_env = e.hsc_unit_env {ue_home_unit_graph = unitEnv_union mergeHugs hug e.hsc_unit_env.ue_home_unit_graph}}
+
+-- | Compile a single module using 'compileHpt' roughly like it happens when Buck requests it.
+-- If the module's home unit has not been encountered before, simulate the metadata step by doctoring the session,
+-- running downsweep, and merging the resulting module graph into the persisted state.
+makeModule :: Conf -> [UnitConf] -> UnitMod -> StateT (Set String) IO ()
+makeModule Conf {..} units umod@UnitMod {unit, src, deps} = do
+  log <- newLog False
+  let env = Env {log, cache, args}
+      envC = env {args = env.args {ghcOptions = env.args.ghcOptions}}
+  firstTime <- state \ seen ->
+    if Set.member unit seen
+    then (False, seen)
+    else (True, Set.insert unit seen)
+  when firstTime do
+    success <- fmap (fromMaybe Nothing) $ liftIO $ withUnitSpecificOptions False env \ env1 specific argv -> do
+      (_, withPackageId) <- liftIO $ readMVar cache <&> \case
+        Cache {hug}
+          | Just h <- hug
+          -> fmap dummyLocation <$> adaptHp h (unLoc <$> argv)
+          | otherwise
+          -> (mempty, argv)
+      flip (withGhcInSession env1) (withPackageId ++ fmap dummyLocation dbsM) \ _ -> do
+        dbg ""
+        dbg (">>> metadata for " ++ unit)
+        modifySession $ hscUpdateFlags \ d -> d {ghcMode = MkDepend}
+        initializeSessionPlugins
+        graph <- loadModuleGraph env1 umod specific
+        pure (Just (Just graph))
+    when debugState do
+      liftIO $ readMVar cache >>= \ Cache {..} -> for_ moduleGraph \ mg ->
+        dbgp (text "updated module graph:" <+> showModGraph mg)
+    void $ case success of
+      Just module_graph ->
+        liftIO $ withGhcMhu envC \ specific _ -> do
+          -- TODO while this makes the modules discoverable in @depanal@, it also causes @compileModuleWithDepsInHpt@ to break,
+          -- when it looks up @unit-main:Err@ for some reason.
+          when False do
+            dbg ""
+            dbg (">>> updating Finder")
+            modifySession $ hscUpdateFlags \ d -> d {ghcMode = CompManager}
+            initUnit specific
+            hsc_env <- getSession
+            liftIO $ for_ (mgModSummaries module_graph) \ m -> do
+              when (homeUnitId (hsc_home_unit hsc_env) == moduleUnitId (ms_mod m)) do
+                dbgp ("add module:" <+> ppr (homeUnitId (hsc_home_unit hsc_env)) <+> ppr (ms_mod_name m))
+                void $ addHomeModuleToFinder (hsc_FC hsc_env) (hsc_home_unit hsc_env) (ms_mod_name m) (ms_location m)
+          pure (Just ())
+      Nothing ->
+        liftIO $ throwGhcExceptionIO (ProgramError "Metadata failed")
+  result <- liftIO $ withGhcMhu envC \ specific target -> do
+    dbg ""
+    dbg (">>> compiling " ++ takeFileName target.get)
+    modifySession $ hscUpdateFlags \ d -> d {ghcMode = CompManager}
+    when False do
+      cache' <- liftIO $ readMVar env.cache
+      modifySessionM \ hsc_env -> do
+        hsc_FC <- liftIO $ newFinderCache env.cache cache' target
+        pure hsc_env {hsc_FC}
+    compileModuleWithDepsInHpt specific target
+  when (isNothing result) do
+      liftIO $ throwGhcExceptionIO (ProgramError "Compile failed")
+  dbgs result
+  where
+    args =
+      args0 {
+        ghcOptions = args0.ghcOptions ++ fileOptions
+      }
+
+    fileOptions =
+      [
+        "-i",
+        -- "-i" ++ takeDirectory src,
+        "-this-unit-id",
+        unit,
+        "-o",
+        tmp </> "out" </> (modName ++ ".dyn_o"),
+        "-ohi",
+        tmp </> "out" </> (modName ++ ".dyn_hi"),
+        src
+      ] ++
+      concat [["-package", dep] | dep <- deps]
+
+    modName = dropExtension (takeFileName src)
+
+    dbsM = concat [["-package-db", db]  | UnitConf {db} <- units]
+
+unitMod :: Conf -> [String] -> String -> String -> String -> UnitMod
+unitMod conf deps name unit content =
+  UnitMod {
+    name,
+    src = conf.tmp </> "src" </> unit </> name ++ ".hs",
+    unit,
+    ..
+  }
+
+errContent :: String
+errContent =
+  unlines [
+    "{-# language DerivingStrategies #-}",
+    "module Err where",
+    "num :: Int",
+    "num = 5",
+    "newtype N = N Int deriving newtype Num"
+  ]
+
+bugContent :: String
+bugContent =
+  unlines [
+    -- "{-# options -ddump-tc -ddump-tc-trace #-}",
+    "module Bug where",
+    "import Language.Haskell.TH",
+    "import Language.Haskell.TH.Syntax",
+    "import Err",
+    "bug :: ExpQ",
+    "bug = lift @_ @Int num",
+    "n :: N",
+    "n = 1"
+  ]
+
+main1 :: String
+main1 =
+  unlines
+    [
+      "{-# language TemplateHaskell #-}",
+      "module Main where",
+      "import Bug",
+      "main :: IO ()",
+      "main = print $(bug)"
+    ]
+
+{-
+-- unit-a:
+module Err where
+num :: Int
+num = 5
+-- unit-main:
+module Bug where
+import Err
+bug :: ExpQ
+bug = lift @_ @Int num
+module Main where
+import Bug
+main :: IO ()
+main = print $(bug)
+-}
+
+targets1 :: Conf -> [UnitMod]
+targets1 conf =
+  [
+    unitMod conf [] "Err" "unit-a" errContent,
+    unitMod conf ["unit-a"] "Bug" "unit-main" bugContent,
+    unitMod conf [] "Main" "unit-main" main1
+  ]
+
+useTh :: Bool
+useTh = True
+
+modType1 :: Conf -> [String] -> Char -> Int -> [String] -> UnitMod
+modType1 conf pdeps unitTag n deps =
+  unitMod conf pdeps modName unit content
+  where
+    unit = "unit-" ++ [unitTag]
+    content =
+      unlines $
+        ("module " ++ modName ++ " where") :
+        "import Language.Haskell.TH (ExpQ)" :
+        "import Language.Haskell.TH.Syntax (lift)" :
+        depImports ++
+        if useTh
+        then [
+          binding ++ " :: ExpQ",
+          binding ++ " = lift @_ @Int 5"
+        ]
+        else [
+          binding ++ " :: Int",
+          binding ++ " = 5"
+        ]
+    depImports = ["import " ++ d | d <- deps]
+    modName = toUpper unitTag : show n
+    binding = unitTag : show n
+
+sumTh :: [String] -> String
+sumTh =
+  foldl' (\ z a -> z ++ " + $(" ++ a ++ ")") "0"
+
+modType2 :: Conf -> [String] -> Char -> Int -> [String] -> [String] -> UnitMod
+modType2 conf pdeps unitTag n deps thDeps =
+  unitMod conf (pdeps) modName unit content
+  where
+    unit = "unit-" ++ [unitTag]
+    content =
+      unlines $
+        "{-# language TemplateHaskell #-}" :
+        ("module " ++ modName ++ " where") :
+        depImports ++
+        [
+          binding ++ " :: Int",
+          binding ++ " = " ++ if useTh
+          then sumTh thDeps
+          else "5"
+        ]
+    depImports = ["import " ++ d | d <- deps]
+    modName = toUpper unitTag : show n
+    binding = unitTag : show n
+
+mainContent :: [(Char, Int)] -> String
+mainContent deps =
+  unlines $
+    "{-# language TemplateHaskell #-}" :
+    "module Main where" :
+    ["import " ++ toUpper c : show i | (c, i) <- deps] ++
+    [
+      "import Bug",
+      "main :: IO ()",
+      "main = do",
+      "  if False then print $(bug) else pure ()",
+      "  print (" ++ (if useTh then sumTh names else intercalate " + " ("0" : names)) ++ ")"
+    ]
+  where
+    names = [c : show i | (c, i) <- deps]
+
+targets2 :: Conf -> [UnitMod]
+targets2 conf =
+  [
+    unitMod conf [] "Err" "unit-b" errContent,
+    m1 'b' 1 [],
+    m1 'b' 2 [],
+    m1d ["unit-b"] 'a' 0 ["B2"],
+    modType2 conf ["unit-b"] 'a' 1 ["B1"] ["b1"],
+    m1 'b' 3 [],
+    m1d ["unit-b"] 'a' 2 ["A0", "A1", "B2", "B3"],
+    unitMod conf ["unit-b"] "Bug" "unit-main" bugContent,
+    unitMod conf ["unit-a", "unit-b"] "Main" "unit-main" (mainContent [
+      ('a', 0),
+      ('a', 2),
+      ('b', 1)
+    ])
+  ]
+  where
+    m1 = modType1 conf []
+
+    m1d = modType1 conf
+
+testWorker :: (Conf -> [UnitMod]) -> IO ()
+testWorker mkTargets =
+  withProject (pure . mkTargets) \ conf units targets -> do
+    evalStateT (traverse_ (makeModule conf units) targets) Set.empty
+
+-- | A very simple test consisting of two home units, using a transitive TH dependency across unit boundaries.
+test_compileHpt :: IO ()
+test_compileHpt = testWorker targets1
diff --git a/buck-worker/test/Main.hs b/buck-worker/test/Main.hs
new file mode 100644
index 00000000..b54a1e6a
--- /dev/null
+++ b/buck-worker/test/Main.hs
@@ -0,0 +1,6 @@
+module Main where
+
+import CompileHptTest (test_compileHpt)
+
+main :: IO ()
+main = test_compileHpt
diff --git a/buck-worker/test/TestSetup.hs b/buck-worker/test/TestSetup.hs
new file mode 100644
index 00000000..620d0b2c
--- /dev/null
+++ b/buck-worker/test/TestSetup.hs
@@ -0,0 +1,175 @@
+module TestSetup where
+
+import Control.Concurrent (MVar)
+import Data.Foldable (for_, toList)
+import Data.Functor ((<&>))
+import qualified Data.List.NonEmpty as NonEmpty
+import Data.List.NonEmpty (NonEmpty (..))
+import GHC.Unit (UnitId, stringToUnitId)
+import Internal.Args (Args (..))
+import Internal.Cache (Cache (..), CacheFeatures (..), emptyCacheWith)
+import Internal.Log (dbg)
+import Prelude hiding (log)
+import System.Directory (createDirectoryIfMissing, listDirectory)
+import System.Environment (getEnv)
+import System.FilePath (takeDirectory, (</>))
+import System.IO.Temp (withSystemTempDirectory)
+import System.Process.Typed (proc, runProcess_)
+
+-- | Global configuration for a worker compilation test.
+data Conf =
+  Conf {
+    -- | Root directory of the test in @/tmp@.
+    tmp :: FilePath,
+
+    -- | The worker cache.
+    cache :: MVar Cache,
+
+    -- | The base cli args used for all modules.
+    args0 :: Args,
+
+    -- | The directory containing the GHC(-pkg) binaries at @bin/@ and the settings (topdir) at @lib/ghc-*/lib/@.
+    ghcDir :: FilePath,
+
+    -- | The relative path to the topdir in 'ghcDir', with the version number spelled out.
+    libPath :: FilePath
+  }
+
+-- | Config for a single test module.
+data UnitMod =
+  UnitMod {
+    -- | Module name.
+    name :: String,
+
+    -- | Path to the source file.
+    src :: FilePath,
+
+    -- | Home unit to which this module belongs.
+    unit :: String,
+
+    -- | Names of home units on which this module depends.
+    deps :: [String],
+
+    -- | The module's source code.
+    content :: String
+  }
+  deriving stock (Eq, Show)
+
+-- | Config for a single test home unit.
+data UnitConf =
+  UnitConf {
+    -- | Unit ID.
+    uid :: UnitId,
+
+    -- | Path to the dummy package DB created for the metadata step, analogous to what's created by Buck.
+    db :: FilePath,
+
+    -- | The modules belonging to this unit.
+    mods :: NonEmpty UnitMod
+  }
+  deriving stock (Eq)
+
+-- | General CLI args used by each module job.
+baseArgs :: FilePath -> FilePath -> Args
+baseArgs topdir tmp =
+  Args {
+    topdir = Just topdir,
+    workerTargetId = Just "test",
+    env = mempty,
+    binPath = [],
+    tempDir = Just (tmp </> "tmp"),
+    ghcPath = Nothing,
+    ghcOptions = (artifactDir =<< ["o", "hie", "dump"]) ++ [
+      "-fwrite-ide-info",
+      "-no-link",
+      "-dynamic",
+      -- "-fwrite-if-simplified-core",
+      "-fbyte-code-and-object-code",
+      "-fprefer-byte-code",
+      -- "-fpackage-db-byte-code",
+      -- "-shared",
+      "-fPIC",
+      "-osuf",
+      "dyn_o",
+      "-hisuf",
+      "dyn_hi",
+      "-package",
+      "base"
+      -- , "-v"
+      -- , "-ddump-if-trace"
+    ]
+  }
+  where
+    artifactDir a = ["-" ++ a ++ "dir", tmp </> "out"]
+
+-- | A package DB config file for the given unit.
+dbConf :: String -> String -> [UnitMod] -> String
+dbConf srcDir unit mods =
+  unlines $ [
+    "name: " ++ unit,
+    "version: 1.0",
+    "id: " ++ unit,
+    "key: " ++ unit,
+    "import-dirs: " ++ srcDir,
+    "exposed: True",
+    "exposed-modules:"
+  ] ++
+  exposed
+  where
+    exposed = [name | UnitMod {name} <- mods]
+
+-- | Write a fresh package DB without a library to the specified directory, using @ghc-pkg@ from the directory in
+-- 'Conf'.
+createDb :: Conf -> String -> String -> IO String
+createDb conf dir confFile = do
+  dbg ("create db for " ++ confFile ++ " at " ++ db)
+  createDirectoryIfMissing False db
+  runProcess_ (proc ghcPkg ["--package-db", db, "recache"])
+  runProcess_ (proc ghcPkg ["--package-db", db, "register", "--force", confFile])
+  pure db
+  where
+    db = dir </> "package.conf.d"
+    ghcPkg = conf.ghcDir </> "bin/ghc-pkg"
+
+-- | Create a package DB for a set of 'UnitMod' and assemble everything into a 'UnitConf'.
+createDbUnitMod :: Conf -> NonEmpty UnitMod -> IO UnitConf
+createDbUnitMod conf mods@(mod0 :| _) = do
+  let uid = stringToUnitId mod0.unit
+  writeFile confFile (dbConf dir mod0.unit (toList mods))
+  db <- createDb conf dir confFile
+  pure UnitConf {..}
+  where
+    confFile = dir </> (mod0.unit ++ ".conf")
+    dir = takeDirectory mod0.src
+
+-- | Set up an environment with dummy package DBs for the set of modules returned by the first argument, then run the
+-- second argument with the resulting unit configurations.
+withProject ::
+  (Conf -> IO [UnitMod]) ->
+  (Conf -> [UnitConf] -> [UnitMod] -> IO a) ->
+  IO a
+withProject mkTargets use =
+  withSystemTempDirectory "buck-worker-test" \ tmp -> do
+    for_ @[] ["src", "tmp", "out"] \ dir ->
+      createDirectoryIfMissing False (tmp </> dir)
+    cache <- emptyCacheWith CacheFeatures {
+      hpt = True,
+      loader = False,
+      enable = True,
+      names = False,
+      finder = True,
+      eps = False
+    }
+    ghcDir <- getEnv "ghc_dir"
+    libPath <- listDirectory (ghcDir </> "lib") <&> \case
+      [d] -> "lib" </> d </> "lib"
+      ds -> error ("weird GHC lib dir contains /= 1 entries: " ++ show ds)
+    let topdir = ghcDir </> libPath
+        conf = Conf {tmp, cache, args0 = baseArgs topdir tmp, ..}
+    targets <- mkTargets conf
+    for_ targets \ UnitMod {src, content} -> do
+      createDirectoryIfMissing False (takeDirectory src)
+      writeFile src content
+    let unitMods = NonEmpty.groupAllWith (.unit) targets
+    units <- traverse (createDbUnitMod conf) unitMods
+    use conf units targets

From 0f0182c573c1759f4a169762e34c316340983576 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Mon, 7 Apr 2025 15:40:38 -0700
Subject: [PATCH 15/16] Install a new CI process.

Change-Id: Idd9322096769e792958d734ea0806a9d75661278
---
 .github/workflows/test.yaml | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)
 create mode 100644 .github/workflows/test.yaml

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
new file mode 100644
index 00000000..a73f76a9
--- /dev/null
+++ b/.github/workflows/test.yaml
@@ -0,0 +1,25 @@
+name: Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+
+  tests:
+    name: Run tests
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      id-token: write
+    steps:
+      - uses: actions/checkout@v4
+      - uses: DeterminateSystems/nix-installer-action@v16
+        with:
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+      - name: Init Cabal
+        run: nix-shell --run 'cabal update'
+      - name: Run tests
+        run: nix-shell --run 'cabal test buck-worker'

From 090e61ff47e19b839bbd0d265fbf42a8515d8bc7 Mon Sep 17 00:00:00 2001
From: Ian-Woo Kim <ianwookim@gmail.com>
Date: Fri, 11 Apr 2025 10:00:26 -0700
Subject: [PATCH 16/16] fix merge mistake

Change-Id: I9218b255a7bbfa650926ac6f6145c32285abe2ea
---
 buck-worker/Main.hs | 22 ----------------------
 1 file changed, 22 deletions(-)

diff --git a/buck-worker/Main.hs b/buck-worker/Main.hs
index 575735cd..35379ed6 100644
--- a/buck-worker/Main.hs
+++ b/buck-worker/Main.hs
@@ -88,7 +88,6 @@ compileAndReadAbiHash ::
   (Target -> Ghc (Maybe ModuleArtifacts)) ->
   Chan (Proto Instr.Event) ->
   BuckArgs ->
-  [String] ->
   Target ->
   Ghc (Maybe CompileResult)
 compileAndReadAbiHash ghcMode compile instrChan args target = do
@@ -107,27 +106,6 @@ compileAndReadAbiHash ghcMode compile instrChan args target = do
         path <- args.abiOut
         Just AbiHash {path, hash = showAbiHash hsc_env artifacts.iface}
     pure CompileResult {artifacts, abiHash}
-  where
-    (ghcMode, compile) = case mode of
-      WorkerOneshotMode -> (OneShot, compileModuleWithDepsInEps)
-      WorkerMakeMode -> (CompManager, compileModuleWithDepsInHpt specific)
-
--- | Process a worker request based on the operational mode specified in the request arguments, either compiling a
--- single module for 'ModeCompile' (@-c@), or computing and writing the module graph to a JSON file for 'ModeMetadata'
--- (@-M@).
-dispatch ::
-  WorkerMode ->
-  Chan (Proto Instr.Event) ->
-  Env ->
-  BuckArgs ->
-  IO Int32
-dispatch workerMode instrChan env args =
-  case args.mode of
-    Just ModeCompile -> do
-      result <- withGhcMhu env (compileAndReadAbiHash workerMode instrChan args)
-      writeResult args result
-    Just m -> error ("worker: mode not implemented: " ++ show m)
-    Nothing -> error "worker: no mode specified"
 
 -- | Process a worker request based on the operational mode specified in the request arguments, either compiling a
 -- single module for 'ModeCompile' (@-c@), or computing and writing the module graph to a JSON file for 'ModeMetadata'