content: {-
    Copyright (c) Meta Platforms, Inc. and affiliates.
    All rights reserved.
  
    This source code is licensed under the BSD-style license found in the
    LICENSE file in the root directory of this source tree.
  -}
  
  {-# LANGUAGE TypeApplications #-}
  
  module Glean.Glass.XRefs
   (
      GenXRef(..),
      XRef,
      XlangXRef,
      resolveEntitiesRange,
      buildGenXRefs,
      fetchCxxIdlXRefs
    ) where
  
  import Data.Map.Strict ( Map )
  
  import qualified Data.Map as Map
  import Data.Bifunctor ( bimap, Bifunctor (first) )
  import Data.List ( foldl' )
  import Data.List.NonEmpty ( NonEmpty(..) )
  import Util.List ( uniq )
  
  import qualified Glean
  import Glean.Angle
  import Glean.Util.ToAngle ( ToAngleFull(toAngleFull), Normalize(normalize) )
  import Glean.Haxl.Repos as Glean ( RepoHaxl )
  import Glean.Glass.Range ( rangeSpanToLocationRange )
  import Glean.Glass.Utils
  import Glean.Glass.Types ( LocationRange, RepoName(..) )
  import qualified Glean.Schema.CodemarkupTypes.Types as Code
  import qualified Glean.Schema.Codemarkup.Types as Code
  import qualified Glean.Schema.Code.Types as Code
  import qualified Glean.Schema.Cxx1.Types as Cxx
  import qualified Glean.Schema.CodemarkupCxx.Types as CodeCxx
  import qualified Glean.Schema.Src.Types as Src
  import Glean.Glass.Query (symbolToEntity)
  import Data.Maybe (mapMaybe)
  import Control.Monad.Extra (mapMaybeM)
  
  type XRef = (Code.XRefLocation, Code.Entity)
  type XlangXRef = (Code.RangeSpan, Either Code.IdlEntity Code.SymbolId)
  data GenXRef = PlainXRef XRef | XlangXRef XlangXRef
  
  -- Turn a Codemarkup "genericEntity" (plain or xlang) into
  -- the corresponding Glass datastructure
  buildGenXRefs :: Code.GenericEntity -> Maybe GenXRef
  buildGenXRefs genEntity = case genEntity of
    Code.GenericEntity_xlangEntity (Code.GenericEntity_xlangEntity_ source entity)
      -> Just $ XlangXRef (source, Left entity)
    Code.GenericEntity_xlangSymbol (Code.GenericEntity_xlangSymbol_ source symbol)
      -> Just $ XlangXRef (source, Right symbol)
    Code.GenericEntity_plainEntity (Code.GenericEntity_plainEntity_ xref entity)
      -> Just $ PlainXRef (xref, entity)
    Code.GenericEntity_EMPTY -> Nothing
  
  -- | First fetch the entities for a given list of symbols.
  --  Then annotate all the entities with the
  --  LocationRange/File of the definitions of these entities.
  --
  --  Note that the rangespans from the parameters are spans for the
  --  references to the entities, not the spans of the entity definitions,
  --  but we want to preserve the correspondence.
  --
  --  This can't be done in Codemarkup for xlang entities as the entity facts
  --  and EntityLocation facts may live in different dbs.
  resolveEntitiesRange
   :: RepoName
   -> [(Code.SymbolId, Code.RangeSpan)]
   -> [(Code.Entity, Code.RangeSpan)]
   -> Glean.RepoHaxl u w
        [((Code.Entity, Code.RangeSpan), (Src.File, LocationRange))]
  resolveEntitiesRange repo symbols ents = do
      entsFromSymbols <- resolveSymbolsToEntities symbols
      resolveEntitiesRanges repo $ uniq $ ents ++ entsFromSymbols
  
  resolveSymbolsToEntities ::
    [(Code.SymbolId, Code.RangeSpan)]
    -> Glean.RepoHaxl u w [(Code.Entity, Code.RangeSpan)]
  resolveSymbolsToEntities symbolRanges = do
    -- We want to query Glean for the entities corresponding to the symbols.
    -- This is more efficiently done in bulk so we need to extract
    -- the symbol list from `symbolRanges`.
    -- To preserve the correspondence between symbols and ranges, we
    -- have to build a map from symbol to entities which we look up at the end.
    let symbols = uniq $ fst <$> symbolRanges
    symbolEnts <- fetchSymbolEntities symbols
    let symbolEntsMap = Map.fromList symbolEnts
    pure $ mapMaybe (mapMaybeFst symbolEntsMap) symbolRanges
    where
      mapMaybeFst :: forall k v a. Ord k => Map k v -> (k, a) -> Maybe (v, a)
      mapMaybeFst map (k, x) = (, x) <$> Map.lookup k map
  
      -- | Query Glean for the entities corresponding to symbols
      fetchSymbolEntities ::
        [Code.SymbolId] -> Glean.RepoHaxl u w [(Code.SymbolId, Code.Entity)]
      fetchSymbolEntities symbols =
        mapMaybeM (fetchDataRecursive . symbolToEntity . toAngleFull) symbols
  
  resolveEntitiesRanges ::
    RepoName
    -> [(Code.Entity, Code.RangeSpan)]
    -> Glean.RepoHaxl u w
        [((Code.Entity, Code.RangeSpan), (Src.File, LocationRange))]
  resolveEntitiesRanges repo entRanges = do
    -- We want to query Glean for the locations corresponding to the entities.
    -- This is more efficiently done in bulk so we need to extract
    -- the entity list from `entRanges`.
    -- To preserve the correspondence between entities and their reference ranges,
    -- we have to build a map from entity to range which we look up at the end.
    entityLocs <- entityLocRange repo (fst <$> entRanges)
    let entityLocsMap = Map.fromList $ first normalize <$> entityLocs
    pure $ mapMaybe (withLoc entityLocsMap) entRanges
    where
      withLoc map entRange@(ent, _) =
        (entRange,) <$> Map.lookup (normalize ent) map
  
      -- | query glean for idl entity locations and
      --   convert those to location ranges
      entityLocRange :: RepoName
        -> [Code.Entity]
        -> RepoHaxl u w [(Code.Entity, (Src.File, LocationRange))]
      entityLocRange reponame refs =
          fetchEntityLocations refs >>= mapM convertSpan
        where
          convertSpan :: (Code.Entity, Code.Location)
            -> RepoHaxl u w (Code.Entity, (Src.File, LocationRange))
          convertSpan (ent, loc) = do
            let file = Code.location_file loc
            range <- rangeSpanToLocationRange reponame file
              (Code.location_location loc)
            return (ent, (file, range))
  
      -- | query glean for idl entity locations
      fetchEntityLocations ::
        [Code.Entity] -> Glean.RepoHaxl u w [(Code.Entity, Code.Location)]
      fetchEntityLocations ents = do
        -- careful to not rely on fact ids in this query
        let angleEntities = toAngleFull <$> ents
        case angleEntities of
          [] -> return []
          hd : tl -> fst <$>
            searchRecursiveWithLimit Nothing (declarationLocation (hd :| tl))
        where
          declarationLocation
            :: NonEmpty (Angle Code.Entity)
            -> Angle (Code.Entity, Code.Location)
          declarationLocation (hd :| tl) =
            let or_ents = foldl' (.|) hd tl in
            vars $ \ent loc ->
              tuple (ent, loc) `where_` [
                or_ents .= ent,
                wild .= predicate @Code.EntityLocation (
                  rec $
                    field @"entity" ent $
                    field @"location" loc
                  end)
              ]
  
  --  Cxx has a specific logic: The only supported xlang xrefs are "IDL" based,
  --  and rely on a mapping "generated entity -> idl entity"
  type EntityIdlMap = Map Code.Entity Code.IdlEntity
  
  -- For Cxx, fetch from Glean the "generated entity -> idl entity" map
  -- and returns an idl extractor which extracts the idl xrefs from
  -- the regular ones. Also propagate the "trunc" status.
  fetchCxxIdlXRefs :: Maybe Int
    -> Glean.IdOf Cxx.FileXRefs
    -> Glean.RepoHaxl u w
      (([XRef], Bool) -> ([GenXRef], Bool))
  fetchCxxIdlXRefs mlimit xrefId =
    do
      (map, trunc) <- entityIdlCxxMap mlimit xrefId
      return $ Data.Bifunctor.bimap (extractIdlXRefs map)(trunc ||)
    where
      entityIdlCxxMap
        :: Maybe Int
        -> Glean.IdOf Cxx.FileXRefs
        -> Glean.RepoHaxl u w (EntityIdlMap, Bool)
      entityIdlCxxMap mlimit xrefId = do
        (rows, truncated) <- searchRecursiveWithLimit mlimit
          (cxxFileEntityIdl xrefId)
        return (Map.fromList rows, truncated)
        where
          cxxFileEntityIdl
            :: Glean.IdOf Cxx.FileXRefs
            -> Angle (Code.Entity, Code.IdlEntity)
          cxxFileEntityIdl xrefId =
            vars $ \(ent :: Angle Code.Entity)
                (idl :: Angle Code.IdlEntity) ->
              tuple (ent, idl) `where_` [
                wild .= predicate @CodeCxx.CxxFileEntityIdl (
                  rec $
                    field @"trace" (asPredicate (factId xrefId)) $
                    field @"ent" ent $
                    field @"idlEnt" idl
                  end)
                ]
  
      -- | extract idl xrefs from the regular ones
      extractIdlXRefs
        :: EntityIdlMap -> [(Code.XRefLocation, Code.Entity)] -> [GenXRef]
      extractIdlXRefs entityIdlMap xRefs =
        let pred (loc, ent) = case Map.lookup ent entityIdlMap of
              Just idl ->
                [PlainXRef (loc, ent),
                  XlangXRef (Code.xRefLocation_source loc, Left idl)]
              _ -> [PlainXRef (loc, ent)] in
        concatMap pred xRefs
  