content: -- | A 'HomeUnitGraph' (HUG) collects information about all the home units.
  -- Crucially, each node in a 'HomeUnitGraph' includes a 'HomePackageTable'.
  --
  -- Often, we don't want to query just a single 'HomePackageTable', but rather all
  -- 'HomePackageTable's of all home units.
  --
  -- This module is responsible for maintaining this bridge between querying all
  -- home units vs querying the home package table directly. Think 'lookupHug' vs
  -- 'lookupHpt', 'hugAllInstances' vs 'hptAllInstances', where the @hug@ version
  -- replies with information from all home units, and the @hpt@ version with
  -- information pertaining to a single home unit.
  --
  -- Meant to be imported qualified as @HUG@.
  -- Example usage:
  --
  -- @
  -- import GHC.Unit.Home.Graph (HomeUnitGraph, HomeUnitEnv)
  -- import qualified GHC.Unit.Home.Graph as HUG
  -- usage = ... HUG.insertHug hug uid modname modinfo ...
  -- @
  module GHC.Unit.Home.Graph
    ( HomeUnitGraph
    , HomeUnitEnv(..)
    , mkHomeUnitEnv
  
    -- * Operations
    , addHomeModInfoToHug
    , restrictHug
    , renameUnitId
    , allUnits
    , updateUnitFlags
  
    -- ** Lookups
    , lookupHug
    , lookupHugByModule
    , lookupHugUnit
    , lookupHugUnitId
    , lookupAllHug
    , memberHugUnit
    , memberHugUnitId
    -- ** Reachability
    , transitiveHomeDeps
  
    -- * Very important queries
    , allInstances
    , allFamInstances
    , allAnns
    , allCompleteSigs
  
    -- * Utilities
    , hugSCCs
    , hugFromList
  
    -- ** Printing
    , pprHomeUnitGraph
    , pprHomeUnitEnv
  
    -- * Auxiliary internal structure
    , UnitEnvGraph(..)
    , unitEnv_lookup_maybe
    , unitEnv_foldWithKey
    , unitEnv_singleton
    , unitEnv_adjust
    , unitEnv_keys
    , unitEnv_insert
    , unitEnv_new
    , unitEnv_lookup
    , unitEnv_traverseWithKey
    , unitEnv_assocs
    ) where
  
  import GHC.Prelude
  
  import GHC.Driver.DynFlags
  import GHC.Unit.Home
  import GHC.Unit.Home.ModInfo
  import GHC.Unit.Home.PackageTable
  import GHC.Unit.Module
  import GHC.Unit.Module.ModIface
  import GHC.Unit.State
  import GHC.Utils.Monad (mapMaybeM)
  import GHC.Utils.Outputable
  import GHC.Utils.Panic
  
  import GHC.Core.FamInstEnv
  
  import Data.Map.Strict (Map)
  import qualified Data.Map.Strict as Map
  import qualified Data.Set as Set
  import GHC.Data.Maybe
  import GHC.Data.Graph.Directed
  
  import GHC.Types.Annotations
  import GHC.Types.CompleteMatch
  import GHC.Core.InstEnv
  
  
  -- | Get all 'CompleteMatches' (arising from COMPLETE pragmas) present across
  -- all home units.
  allCompleteSigs :: HomeUnitGraph -> IO CompleteMatches
  allCompleteSigs hug = foldr go (pure []) hug where
    go hue = liftA2 (++) (hptCompleteSigs (homeUnitEnv_hpt hue))
  
  -- | Find all the instance declarations (of classes and families) from
  -- the Home Package Table filtered by the provided predicate function.
  -- Used in @tcRnImports@, to select the instances that are in the
  -- transitive closure of imports from the currently compiled module.
  allInstances :: HomeUnitGraph -> IO (InstEnv, [FamInst])
  allInstances hug = foldr go (pure (emptyInstEnv, [])) hug where
    go hue = liftA2 (\(a,b) (a',b') -> (a `unionInstEnv` a', b ++ b'))
                    (hptAllInstances (homeUnitEnv_hpt hue))
  
  allFamInstances :: HomeUnitGraph -> IO (ModuleEnv FamInstEnv)
  allFamInstances hug = foldr go (pure emptyModuleEnv) hug where
    go hue = liftA2 plusModuleEnv (hptAllFamInstances (homeUnitEnv_hpt hue))
  
  allAnns :: HomeUnitGraph -> IO AnnEnv
  allAnns hug = foldr go (pure emptyAnnEnv) hug where
    go hue = liftA2 plusAnnEnv (hptAllAnnotations (homeUnitEnv_hpt hue))
  
  --------------------------------------------------------------------------------
  -- HomeUnitGraph (HUG)
  --------------------------------------------------------------------------------
  
  type HomeUnitGraph = UnitEnvGraph HomeUnitEnv
  
  data HomeUnitEnv = HomeUnitEnv
    { homeUnitEnv_units     :: !UnitState
        -- ^ External units
  
    , homeUnitEnv_unit_dbs :: !(Maybe [UnitDatabase UnitId])
        -- ^ Stack of unit databases for the target platform.
        --
        -- This field is populated with the result of `initUnits`.
        --
        -- 'Nothing' means the databases have never been read from disk.
        --
        -- Usually we don't reload the databases from disk if they are
        -- cached, even if the database flags changed!
  
    , homeUnitEnv_dflags :: DynFlags
      -- ^ The dynamic flag settings
    , homeUnitEnv_hpt :: HomePackageTable
      -- ^ The home package table describes already-compiled
      -- home-package modules, /excluding/ the module we
      -- are compiling right now.
      -- (In one-shot mode the current module is the only
      -- home-package module, so homeUnitEnv_hpt is empty.  All other
      -- modules count as \"external-package\" modules.
      -- However, even in GHCi mode, hi-boot interfaces are
      -- demand-loaded into the external-package table.)
      --
      -- 'homeUnitEnv_hpt' is not mutable because we only demand-load
      -- external packages; the home package is eagerly
      -- loaded, module by module, by the compilation manager.
      --
      -- The HPT may contain modules compiled earlier by @--make@
      -- but not actually below the current module in the dependency
      -- graph.
      --
      -- (This changes a previous invariant: changed Jan 05.)
  
    , homeUnitEnv_home_unit :: !(Maybe HomeUnit)
      -- ^ Home-unit
    }
  
  mkHomeUnitEnv :: UnitState -> Maybe [UnitDatabase UnitId] -> DynFlags -> HomePackageTable -> Maybe HomeUnit -> HomeUnitEnv
  mkHomeUnitEnv us dbs dflags hpt home_unit = HomeUnitEnv
    { homeUnitEnv_units = us
    , homeUnitEnv_unit_dbs = dbs
    , homeUnitEnv_dflags = dflags
    , homeUnitEnv_hpt = hpt
    , homeUnitEnv_home_unit = home_unit
    }
  
  --------------------------------------------------------------------------------
  -- * Operations on HUG
  --------------------------------------------------------------------------------
  
  -- | Add an entry to the 'HomePackageTable' under the unit of that entry.
  addHomeModInfoToHug :: HomeModInfo -> HomeUnitGraph -> IO ()
  addHomeModInfoToHug hmi hug =
    case unitEnv_lookup_maybe hmi_unit hug of
      Nothing -> pprPanic "addHomeInfoToHug" (ppr hmi_mod)
      Just hue -> do
        addHomeModInfoToHpt hmi (homeUnitEnv_hpt hue)
    where
      hmi_mod :: Module
      hmi_mod  = mi_module (hm_iface hmi)
      hmi_unit = toUnitId (moduleUnit hmi_mod)
  
  -- | 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.
  restrictHug :: [(UnitId, [HomeModInfo])] -> HomeUnitGraph -> IO ()
  restrictHug deps hug = unitEnv_foldWithKey (\k uid hue -> restrict_one uid hue >> k) (return ()) hug
    where
      deps_map = Map.fromList deps
      restrict_one uid hue  =
        restrictHpt (homeUnitEnv_hpt hue) (Map.findWithDefault [] uid deps_map)
  
  -- | Rename a unit id in the 'HomeUnitGraph'
  --
  -- @'renameUnitId' oldUnit newUnit hug@, if @oldUnit@ is not found in @hug@, returns 'Nothing'.
  -- If it exists, the result maps @newUnit@ to the 'HomeUnitEnv' of the
  -- @oldUnit@ (and @oldUnit@ is removed from @hug@)
  renameUnitId :: UnitId -> UnitId -> HomeUnitGraph -> Maybe HomeUnitGraph
  renameUnitId oldUnit newUnit hug = case unitEnv_lookup_maybe oldUnit hug of
    Nothing -> Nothing
    Just oldHue -> pure $
      unitEnv_insert newUnit oldHue $
      unitEnv_delete oldUnit hug
  
  -- | Retrieve all 'UnitId's of units in the 'HomeUnitGraph'.
  allUnits :: HomeUnitGraph -> Set.Set UnitId
  allUnits = unitEnv_keys
  
  -- | Set the 'DynFlags' of the 'HomeUnitEnv' for unit in the 'HomeModuleGraph'
  updateUnitFlags :: UnitId -> (DynFlags -> DynFlags) -> HomeUnitGraph -> HomeUnitGraph
  updateUnitFlags uid f = unitEnv_adjust update uid
    where
      update hue = hue { homeUnitEnv_dflags = f (homeUnitEnv_dflags hue) }
  
  --------------------------------------------------------------------------------
  -- ** Reachability
  --------------------------------------------------------------------------------
  
  -- | Compute the transitive closure of a unit in the 'HomeUnitGraph'.
  -- If the argument unit is not present in the graph returns Nothing.
  transitiveHomeDeps :: UnitId -> HomeUnitGraph -> Maybe [UnitId]
  transitiveHomeDeps uid hug = case lookupHugUnitId uid hug of
    Nothing -> Nothing
    Just hue -> Just $
      Set.toList (loop (Set.singleton uid) (homeUnitDepends (homeUnitEnv_units hue)))
      where
        loop acc [] = acc
        loop acc (uid:uids)
          | uid `Set.member` acc = loop acc uids
          | otherwise =
            let hue = homeUnitDepends
                      . homeUnitEnv_units
                      . expectJust
                      $ lookupHugUnitId uid hug
            in loop (Set.insert uid acc) (hue ++ uids)
  
  --------------------------------------------------------------------------------
  -- ** Lookups
  --------------------------------------------------------------------------------
  
  -- | Lookup the 'HomeModInfo' of a 'Module' in the 'HomeUnitGraph' given its
  -- 'UnitId' and 'ModuleName' (via the 'HomePackageTable' of the corresponding unit)
  lookupHug :: HomeUnitGraph -> UnitId -> ModuleName -> IO (Maybe HomeModInfo)
  lookupHug hug uid mod = do
    case unitEnv_lookup_maybe uid hug of
      Nothing -> pure Nothing
      Just hue -> lookupHpt (homeUnitEnv_hpt hue) mod
  
  -- | Lookup the 'HomeModInfo' of a 'Module' in the 'HomeUnitGraph' (via the 'HomePackageTable' of the corresponding unit)
  lookupHugByModule :: Module -> HomeUnitGraph -> IO (Maybe HomeModInfo)
  lookupHugByModule mod hug =
    case lookupHugUnit (moduleUnit mod) hug of
      Nothing -> pure Nothing
      Just env -> lookupHptByModule (homeUnitEnv_hpt env) mod
  
  -- | Lookup all 'HomeModInfo' that have the same 'ModuleName' as the given 'ModuleName'.
  -- 'ModuleName's are not unique in the case of multiple home units, so there can be
  -- more than one possible 'HomeModInfo'.
  --
  -- You should always prefer 'lookupHug' and 'lookupHugByModule' when possible.
  lookupAllHug :: HomeUnitGraph -> ModuleName -> IO [HomeModInfo]
  lookupAllHug hug mod = mapMaybeM (\uid -> lookupHug hug uid mod) (Set.toList $ unitEnv_keys hug)
  
  -- | Lookup a 'HomeUnitEnv' by 'UnitId' in a 'HomeUnitGraph'
  lookupHugUnitId :: UnitId -> HomeUnitGraph -> Maybe HomeUnitEnv
  lookupHugUnitId = unitEnv_lookup_maybe
  
  -- | Check whether the 'UnitId' is present in the 'HomeUnitGraph'
  memberHugUnitId :: UnitId -> HomeUnitGraph -> Bool
  memberHugUnitId u = isJust . lookupHugUnitId u
  
  -- | Lookup up the 'HomeUnitEnv' by the 'Unit' in the 'HomeUnitGraph'.
  -- If the 'Unit' can be turned into a 'UnitId', we behave identical to 'lookupHugUnitId'.
  --
  -- A 'HoleUnit' is never part of the 'HomeUnitGraph', only instantiated 'Unit's
  lookupHugUnit :: Unit -> HomeUnitGraph -> Maybe HomeUnitEnv
  lookupHugUnit unit hug =
    if isHoleUnit unit
      then Nothing
      else lookupHugUnitId (toUnitId unit) hug
  
  -- | Check whether the 'Unit' is present in the 'HomeUnitGraph'
  --
  -- A 'HoleUnit' is never part of the 'HomeUnitGraph', only instantiated 'Unit's
  memberHugUnit :: Unit -> HomeUnitGraph -> Bool
  memberHugUnit u = isJust . lookupHugUnit u
  
  --------------------------------------------------------------------------------
  -- * Internal representation map
  --------------------------------------------------------------------------------
  -- Note: we purposefully do not export functions like "elems" to maintain a
  -- good clean interface with the HUG.
  
  type UnitEnvGraphKey = UnitId
  
  newtype UnitEnvGraph v = UnitEnvGraph
    { unitEnv_graph :: Map UnitEnvGraphKey v
    } deriving (Functor, Foldable, Traversable)
  
  unitEnv_new :: Map UnitEnvGraphKey v -> UnitEnvGraph v
  unitEnv_new m =
    UnitEnvGraph
      { unitEnv_graph = m
      }
  
  unitEnv_insert :: UnitEnvGraphKey -> v -> UnitEnvGraph v -> UnitEnvGraph v
  unitEnv_insert unitId env unitEnv = unitEnv
    { unitEnv_graph = Map.insert unitId env (unitEnv_graph unitEnv)
    }
  
  unitEnv_delete :: UnitEnvGraphKey -> UnitEnvGraph v -> UnitEnvGraph v
  unitEnv_delete uid unitEnv =
      unitEnv
        { unitEnv_graph = Map.delete uid (unitEnv_graph unitEnv)
        }
  
  unitEnv_adjust :: (v -> v) -> UnitEnvGraphKey -> UnitEnvGraph v -> UnitEnvGraph v
  unitEnv_adjust f uid unitEnv = unitEnv
    { unitEnv_graph = Map.adjust f uid (unitEnv_graph unitEnv)
    }
  
  unitEnv_singleton :: UnitEnvGraphKey -> v -> UnitEnvGraph v
  unitEnv_singleton active m = UnitEnvGraph
    { unitEnv_graph = Map.singleton active m
    }
  
  unitEnv_lookup_maybe :: UnitEnvGraphKey -> UnitEnvGraph v -> Maybe v
  unitEnv_lookup_maybe u env = Map.lookup u (unitEnv_graph env)
  
  unitEnv_keys :: UnitEnvGraph v -> Set.Set UnitEnvGraphKey
  unitEnv_keys env = Map.keysSet (unitEnv_graph env)
  
  unitEnv_foldWithKey :: (b -> UnitEnvGraphKey -> a -> b) -> b -> UnitEnvGraph a -> b
  unitEnv_foldWithKey f z (UnitEnvGraph g)= Map.foldlWithKey' f z g
  
  unitEnv_lookup :: UnitEnvGraphKey -> UnitEnvGraph v -> v
  unitEnv_lookup u env = expectJust $ unitEnv_lookup_maybe u env
  
  unitEnv_traverseWithKey :: Applicative f => (UnitEnvGraphKey -> a -> f b) -> UnitEnvGraph a -> f (UnitEnvGraph b)
  unitEnv_traverseWithKey f unitEnv =
    UnitEnvGraph <$> Map.traverseWithKey f (unitEnv_graph unitEnv)
  
  unitEnv_assocs :: UnitEnvGraph a -> [(UnitEnvGraphKey, a)]
  unitEnv_assocs (UnitEnvGraph x) = Map.assocs x
  
  --------------------------------------------------------------------------------
  -- * Utilities
  --------------------------------------------------------------------------------
  
  hugSCCs :: HomeUnitGraph -> [SCC UnitId]
  hugSCCs hug = sccs where
    mkNode :: (UnitId, HomeUnitEnv) -> Node UnitId UnitId
    mkNode (uid, hue) = DigraphNode uid uid (homeUnitDepends (homeUnitEnv_units hue))
    nodes = map mkNode (Map.toList $ unitEnv_graph hug)
  
    sccs = stronglyConnCompFromEdgedVerticesOrd nodes
  
  hugFromList :: [(UnitId, HomeUnitEnv)] -> HomeUnitGraph
  hugFromList = UnitEnvGraph . Map.fromList
  
  pprHomeUnitGraph :: HomeUnitGraph -> IO SDoc
  pprHomeUnitGraph unitEnv = do
    docs <- mapM (\(k, v) -> pprHomeUnitEnv k v) $ Map.assocs $ unitEnv_graph unitEnv
    return $ vcat docs
  
  pprHomeUnitEnv :: UnitId -> HomeUnitEnv -> IO SDoc
  pprHomeUnitEnv uid env = do
    hptDoc <- pprHPT $ homeUnitEnv_hpt env
    return $
      ppr uid <+> text "(flags:" <+> ppr (homeUnitId_ $ homeUnitEnv_dflags env) <> text "," <+> ppr (fmap homeUnitId $ homeUnitEnv_home_unit env) <> text ")" <+> text "->"
      $$ nest 4 hptDoc
  
  