content: {-# LANGUAGE LambdaCase #-}
  -- | The 'HomePackageTable' (HPT) contains information about all modules that are part
  -- of a home package. At its core, the information for each module is a
  -- 'HomeModInfo'.
  --
  -- During upsweep, the HPT is a monotonically increasing data structure: it
  -- only ever gets extended by inserting modules which are loaded and for which
  -- we discover the information required to construct a 'ModInfo'.
  --
  -- There should only ever exist one single HPT for any given home unit. It's
  -- crucial we don't accidentally leak HPTs (e.g. by filtering it, which used to
  -- happen -- #25511), so the HPT is mutable and only its reference should be shared.
  -- This is alright because the modules don't change throughout compilation.
  --
  -- :::WARNING:::
  -- If you intend to change this interface, consider carefully whether you are
  -- exposing memory-leak footguns which may end up being misused in the compiler
  -- eventually. For instance, if you really, really, end up needing a way to take
  -- a snapshot of the IORef (think: do you really need to?), at least make
  -- obvious in the name like `snapshotCopyHpt`.
  --
  -- Or, do you really need a function to traverse all modules in the HPT? It is
  -- often better to keep the computation internal to this module, such as in
  -- 'hptCollectObjects'...
  module GHC.Unit.Home.PackageTable
    (
      HomePackageTable(..)
    , emptyHomePackageTable
  
      -- * Lookups in the HPT
    , lookupHpt
    , lookupHptByModule
  
      -- * Extending the HPT
    , addHomeModInfoToHpt
    , addHomeModInfosToHpt
  
      -- * Restrict the HPT
    , restrictHpt
  
      -- * Queries about home modules
    , hptCompleteSigs
    , hptAllInstances
    , hptAllFamInstances
    , hptAllAnnotations
  
      -- ** More Traversal-based queries
    , hptCollectDependencies
    , hptCollectObjects
    , hptCollectModules
  
      -- ** Memory dangerous queries
    , concatHpt
  
      -- * Utilities
    , pprHPT
  
      -- * Internals
      --
      -- | These provide access to the internals of the HomePackageTable to
      -- facilitate existing workflows that used the previous API. For instance,
      -- if you were listing out all elements or merging, you can keep doing so by reading
      -- the internal IO ref and then using the moduleenv contents directly.
      --
      -- In GHC itself these should be avoided, and other uses should justify why
      -- it is not sufficient to go through the intended insert-only API.
    , hptInternalTableRef
    , hptInternalTableFromRef
  
      -- * Legacy API
      --
      -- | This API is deprecated and meant to be removed.
    , addToHpt
    , addListToHpt
    ) where
  
  import GHC.Prelude
  import GHC.Data.Maybe
  
  import Data.IORef
  import Control.Monad ((<$!>))
  import qualified Data.Set as Set
  
  import GHC.Core.FamInstEnv
  import GHC.Core.InstEnv
  import GHC.Linker.Types
  import GHC.Types.Annotations
  import GHC.Types.CompleteMatch
  import GHC.Types.Unique.DFM
  import GHC.Unit.Home.ModInfo
  import GHC.Unit.Module
  import GHC.Unit.Module.Deps
  import GHC.Unit.Module.ModDetails
  import GHC.Unit.Module.ModIface
  import GHC.Utils.Outputable
  import GHC.Types.Unique (getUnique, getKey)
  import qualified GHC.Data.Word64Set as W64
  
  -- | Helps us find information about modules in the home package
  newtype HomePackageTable = HPT {
  
      table :: IORef (DModuleNameEnv HomeModInfo)
      -- ^ Domain = modules in this home unit
      --
      -- This is an IORef because we want to avoid leaking HPTs (see the particularly bad #25511).
      -- Moreover, the HPT invariant allows mutability in this table without compromising thread safety or soundness.
      -- To recall:
      --   A query to the HPT should depend only on data relevant to that query, such that
      --   there being more or less unrelated entries in the HPT does not influence the result in any way.
      --
      -- Note that the HPT increases monotonically, except at certain barrier
      -- points like when 'restrictHpt' is called. At these barriers, it is safe
      -- to temporarily violate the HPT monotonicity.
      --
      -- The elements of this table may be updated (e.g. on rehydration).
    }
  
  -- | Create a new 'HomePackageTable'.
  --
  -- Be careful not to share it across e.g. different units, since it uses a
  -- mutable variable under the hood to keep the monotonically increasing list of
  -- loaded modules.
  emptyHomePackageTable :: IO HomePackageTable
  -- romes:todo: use a MutableArray directly?
  emptyHomePackageTable = do
    table <- newIORef emptyUDFM
    return HPT{table}
  
  --------------------------------------------------------------------------------
  -- * Lookups in the HPT
  --------------------------------------------------------------------------------
  
  -- | Lookup the 'HomeModInfo' of a module in the HPT, given its name.
  lookupHpt :: HomePackageTable -> ModuleName -> IO (Maybe HomeModInfo)
  lookupHpt HPT{table=hpt} mn = (`lookupUDFM` mn) <$!> readIORef hpt
  
  -- | Lookup the 'HomeModInfo' of a 'Module' in the HPT.
  lookupHptByModule :: HomePackageTable -> Module -> IO (Maybe HomeModInfo)
  lookupHptByModule hpt mod
    = -- The HPT is indexed by ModuleName, not Module,
      -- we must check for a hit on the right Module
      lookupHpt hpt (moduleName mod) >>= pure . \case
        Just hm | mi_module (hm_iface hm) == mod -> Just hm
        _otherwise                               -> Nothing
  
  --------------------------------------------------------------------------------
  -- * Extending the HPT
  --------------------------------------------------------------------------------
  
  -- | Add a new module to the HPT.
  --
  -- An HPT is a monotonically increasing data structure, holding information about loaded modules in a package.
  -- This is the main function by which the HPT is extended or updated.
  --
  -- When the module of the inserted 'HomeModInfo' does not exist, a new entry in
  -- the HPT is created for that module name.
  -- When the module already has an entry, inserting a new one entry in the HPT
  -- will always overwrite the existing entry for that module.
  addHomeModInfoToHpt :: HomeModInfo -> HomePackageTable -> IO ()
  addHomeModInfoToHpt hmi hpt = addToHpt hpt (moduleName (mi_module (hm_iface hmi))) hmi
  
  {-# DEPRECATED addToHpt "Deprecated in favour of 'addHomeModInfoToHpt', as the module at which a 'HomeModInfo' is inserted should always be derived from the 'HomeModInfo' itself." #-}
  -- After deprecation cycle, move `addToHpt` to a `where` clause inside `addHomeModInfoToHpt`.
  addToHpt :: HomePackageTable -> ModuleName -> HomeModInfo -> IO ()
  addToHpt HPT{table=hptr} mn hmi = do
    atomicModifyIORef' hptr (\hpt -> (addToUDFM hpt mn hmi, ()))
    -- If the key already existed in the map, this insertion is overwriting
    -- the HMI of a previously loaded module (likely in rehydration).
  
  -- | 'addHomeModInfoToHpt' for multiple module infos.
  addHomeModInfosToHpt :: HomePackageTable -> [HomeModInfo] -> IO ()
  addHomeModInfosToHpt hpt = mapM_ (flip addHomeModInfoToHpt hpt)
  
  -- | Thin each HPT variable to only contain keys from the given dependencies.
  -- This is used at the end of upsweep to make sure that only completely successfully loaded
  -- modules are visible for subsequent operations.
  --
  -- This is an exception to the invariant of the HPT -- that it grows
  -- monotonically, never removing entries -- which is safe as long as it is only
  -- called at barrier points, such as the end of upsweep, when all threads are
  -- done and we want to clean up failed entries.
  restrictHpt :: HomePackageTable -> [HomeModInfo] -> IO ()
  restrictHpt HPT{table=hptr} hmis =
    let key_set = map (getKey . getUnique . hmi_mod) hmis
        hmi_mod hmi = moduleName (mi_module (hm_iface hmi))
    in atomicModifyIORef' hptr (\hpt -> (udfmRestrictKeysSet hpt (W64.fromList key_set), ()))
  
  {-# DEPRECATED addListToHpt "Deprecated in favour of 'addHomeModInfosToHpt', as the module at which a 'HomeModInfo' is inserted should always be derived from the 'HomeModInfo' itself." #-}
  -- After deprecation cycle, remove.
  addListToHpt :: HomePackageTable -> [(ModuleName, HomeModInfo)] -> IO ()
  addListToHpt hpt = mapM_ (uncurry (addToHpt hpt))
  
  ----------------------------------------------------------------------------------
  ---- * Queries
  ----------------------------------------------------------------------------------
  
  -- | Get all 'CompleteMatches' (arising from COMPLETE pragmas) present in all
  -- modules from this unit's HPT.
  hptCompleteSigs :: HomePackageTable -> IO CompleteMatches
  hptCompleteSigs = concatHpt (md_complete_matches . hm_details)
  
  -- | Find all the instance declarations (of classes and families) from this Home Package Table
  hptAllInstances :: HomePackageTable -> IO (InstEnv, [FamInst])
  hptAllInstances hpt = do
    hits <- flip concatHpt hpt $ \mod_info -> do
       let details = hm_details mod_info
       return (md_insts details, md_fam_insts details)
    let (insts, famInsts) = unzip hits
    return (foldl' unionInstEnv emptyInstEnv insts, concat famInsts)
  
  -- | Find all the family instance declarations from the HPT
  hptAllFamInstances :: HomePackageTable -> IO (ModuleEnv FamInstEnv)
  hptAllFamInstances = fmap mkModuleEnv . concatHpt (\hmi -> [(hmiModule hmi, hmiFamInstEnv hmi)])
    where
      hmiModule     = mi_module . hm_iface
      hmiFamInstEnv = extendFamInstEnvList emptyFamInstEnv
                        . md_fam_insts . hm_details
  
  -- | All annotations from the HPT
  hptAllAnnotations :: HomePackageTable -> IO AnnEnv
  hptAllAnnotations = fmap mkAnnEnv . concatHpt (md_anns . hm_details)
  
  
  --------------------------------------------------------------------------------
  -- * Traversal-based queries
  --------------------------------------------------------------------------------
  
  -- | Collect the immediate dependencies of all modules in the HPT into a Set.
  -- The immediate dependencies are given by the iface as @'dep_direct_pkgs' . 'mi_deps'@.
  --
  -- Note: This should be a query on the 'ModuleGraph', since we don't really
  -- ever want to collect *all* dependencies. The current caller of this function
  -- currently takes all dependencies only to then filter them with an ad-hoc transitive closure check.
  -- See #25639
  hptCollectDependencies :: HomePackageTable -> IO (Set.Set (IfaceImportLevel, UnitId))
  hptCollectDependencies HPT{table} = do
    hpt <- readIORef table
    return $
      foldr (Set.union . dep_direct_pkgs . mi_deps . hm_iface) Set.empty hpt
  
  -- | Collect the linkable object of all modules in the HPT.
  -- The linkable objects are given by @'homeModInfoObject'@.
  --
  -- $O(n)$ in the number of modules in the HPT.
  hptCollectObjects :: HomePackageTable -> IO [Linkable]
  hptCollectObjects HPT{table} = do
    hpt <- readIORef table
    return $
      foldr ((:) . expectJust . homeModInfoObject) [] hpt
  
  -- | Collect all module ifaces in the HPT
  --
  -- $O(n)$ in the number of modules in the HPT.
  hptCollectModules :: HomePackageTable -> IO [Module]
  hptCollectModules HPT{table} = do
    hpt <- readIORef table
    return $
      foldr ((:) . mi_module . hm_iface) [] hpt
  
  --------------------------------------------------------------------------------
  -- * Utilities
  --------------------------------------------------------------------------------
  
  -- | Pretty print a 'HomePackageTable'.
  --
  -- Make sure you really do need to print the whole HPT before infusing too much
  -- code with IO.
  --
  -- For instance, in the HUG, it suffices to print the unit-keys present in the
  -- unit map in failed lookups.
  pprHPT :: HomePackageTable -> IO SDoc
  -- A bit arbitrary for now
  pprHPT HPT{table=hptr} = do
    hpt <- readIORef hptr
    return $!
      pprUDFM hpt $ \hms ->
        vcat [ ppr (mi_module (hm_iface hm))
             | hm <- hms ]
  
  ----------------------------------------------------------------------------------
  -- THE TYPE OF FOOTGUNS WE DON'T WANT TO EXPOSE
  ----------------------------------------------------------------------------------
  
  -- eltsHpt :: HomePackageTable -> [HomeModInfo]
  -- filterHpt :: (HomeModInfo -> Bool) -> HomePackageTable -> HomePackageTable
  -- mapHpt :: (HomeModInfo -> HomeModInfo) -> HomePackageTable -> HomePackageTable
  -- delFromHpt :: HomePackageTable -> ModuleName -> HomePackageTable
  -- listToHpt :: [(ModuleName, HomeModInfo)] -> HomePackageTable
  -- listHMIToHpt :: [HomeModInfo] -> HomePackageTable
  
  ----------------------------------------------------------------------------------
  -- Would be fine, but may lead to linearly traversing the HPT unnecessarily
  -- (e.g. `lastLoadedKey` superseded bad usages)
  ----------------------------------------------------------------------------------
  
  -- allHpt :: (HomeModInfo -> Bool) -> HomePackageTable -> Bool
  -- allHpt = allUDFM
  
  -- anyHpt :: (HomeModInfo -> Bool) -> HomePackageTable -> Bool
  -- anyHpt = anyUDFM
  
  ----------------------------------------------------------------------------------
  -- Would be ok to expose this function very /careful/ with the argument function
  ----------------------------------------------------------------------------------
  
  -- | Like @concatMap f . 'eltsHpt'@, but filters out all 'HomeModInfo' for which
  -- @f@ returns the empty list before doing the sort inherent to 'eltsUDFM'.
  --
  -- If this function is ever exposed from the HPT module, make sure the
  -- argument function doesn't introduce leaks.
  concatHpt :: (HomeModInfo -> [a]) -> HomePackageTable -> IO [a]
  concatHpt f HPT{table} = do
    hpt <- readIORef table
    return $ concat . eltsUDFM . mapMaybeUDFM g $ hpt
    where
      g hmi = case f hmi of { [] -> Nothing; as -> Just as }
  
  --------------------------------------------------------------------------------
  -- * Internals (see haddocks!)
  --------------------------------------------------------------------------------
  
  -- | Gets the internal 'IORef' which holds the 'HomeModInfo's of this HPT.
  -- Use with care.
  hptInternalTableRef :: HomePackageTable -> IORef (DModuleNameEnv HomeModInfo)
  hptInternalTableRef = table
  
  -- | Construct a HomePackageTable from the IORef.
  -- Use with care, only if you can really justify going around the intended insert-only API.
  hptInternalTableFromRef :: IORef (DModuleNameEnv HomeModInfo) -> IO HomePackageTable
  hptInternalTableFromRef ref = do
    return HPT {
      table = ref
    }
  
  