Update symbol Id parsing to allow filters and use that to search for references

Summary:
# Doc

https://docs.google.com/document/d/1zlJlGSmDhQxuJ87K1bH59wPJk2IyY5WMQFTEgKA4Ak0/edit?tab=t.0#heading=h.qn82j2uc6gbw

# Stack
- Currently for directory branches like xlformers-branches when finding references, the references for all branches are returned. This is because the **symbol id is the same across branches**.
- In order to fix this in a more general way (that can apply to other directory based branches) we are implementing nhawkes design of adding filters to symbol ID's which can then be used to decide which references to return.
 - With this set, glass will now only issue a new query, where the go-to-definition for the symbol's match. This means that the go-to-defintion for the symbol you searched for must point to the same file to that the references searched for point to.
- When codehub issues a request, we append this filter since the go-to-defintion information is available at load time.

# This diff
- Add a thrift `union` type which describes the filters we will allow on SymbolId.
- Update parsing logic for symbolId to return a list of SymbolFilters.
- Update `findReferenceRanges` in glass to extract the definition file from the symbol filters if it exists. If this exists issue a new query.
- Add a new query which adds constraints onto the EntityReferences predicate such that the filename of the reference's xref must match the filename of the original entities xref. (Both the go-to-definition's must match).

Reviewed By: pepeiborra

Differential Revision: D78165656

fbshipit-source-id: ea1a3697b4e887dc060211224ac58f4fe4eed8d9

Author: aahanaggarwal

glean/glass/Glean/Glass/Handler/Symbols.hs

@@ -109,14 +109,22 @@ import qualified Glean.Glass.Handler.Utils as Utils
 import Glean.Glass.SymbolKind (findSymbolKind)
 import Glean.Glass.Env (Env' (tracer, sourceControl, useSnapshotsForSymbolsList))
 
+extractDefFileFilter :: [SymbolFilter] -> Maybe Text
+extractDefFileFilter [] = Nothing
+extractDefFileFilter (SymbolFilter_definition_file filePath : _) = Just filePath
+extractDefFileFilter (_ : rest) = extractDefFileFilter rest
+
+
 -- | Symbol-based find-references.
 findReferenceRanges
   :: Glass.Env -> SymbolId -> RequestOptions -> IO [LocationRange]
 findReferenceRanges env@Glass.Env{..} sym opts@RequestOptions{..} =
   withSymbol "findReferenceRanges" env opts sym
-    $ \gleanDBs _dbInfo (repo, lang, toks) ->
-      fetchSymbolReferenceRanges env repo lang toks limit
-        GleanBackend{..}
+    $ \gleanDBs _dbInfo (repo, lang, toks, symbolFilters) -> do
+      let filename = extractDefFileFilter symbolFilters
+      (ranges, maybeError) <-
+            fetchSymbolReferenceRanges env repo lang toks limit filename GleanBackend{..}
+      return (ranges, maybeError)
   where
     limit = fmap fromIntegral requestOptions_limit
 
@@ -125,8 +133,11 @@ findReferenceRanges env@Glass.Env{..} sym opts@RequestOptions{..} =
 symbolLocation :: Glass.Env -> SymbolId -> RequestOptions -> IO SymbolLocation
 symbolLocation env@Glass.Env{..} sym opts =
   withSymbol "symbolLocation" env opts sym $
-    \gleanDBs dbInfo (scmRepo, lang, toks) ->
+    \gleanDBs dbInfo (scmRepo, lang, toks, filters) ->
       backendRunHaxl GleanBackend{..} env $ do
+        case filters of
+          [] -> return ()
+          _ -> throwM (ServerException "symbolLocation: filters not supported")
         r <- Search.searchEntityLocation lang toks
         (entity, err) <- case r of
           None t -> throwM (ServerException t)
@@ -150,8 +161,11 @@ describeSymbol
   -> IO SymbolDescription
 describeSymbol env@Glass.Env{..} symId opts =
   withSymbol "describeSymbol" env opts symId $
-    \gleanDBs dbInfo (scmRepo, lang, toks) ->
+    \gleanDBs dbInfo (scmRepo, lang, toks, filters) ->
       backendRunHaxl GleanBackend{..} env $ do
+        case filters of
+          [] -> return ()
+          _ -> throwM (ServerException "describeSymbol: filters not supported")
         r <- Search.searchEntityLocation lang toks
         (first :| rest, err) <- case r of
           None t -> throwM (ServerException t)
@@ -178,8 +192,11 @@ resolveSymbols env@Glass.Env{..} (ResolveSymbolsRequest symIds) opts = do
   results <- forM symIds $ \symId -> do
     -- Try to resolve the symbol, returning any error encountered
     (resolutions, err) <- withSymbol "resolveSymbols" env opts symId $
-      \gleanDBs dbInfo (scmRepo, lang, toks) ->
+      \gleanDBs dbInfo (scmRepo, lang, toks, filters) ->
       backendRunHaxl GleanBackend{..} env $ do
+        case filters of
+          [] -> return ()
+          _ -> throwM (ServerException "resolveSymbols: filters not supported")
         r <- Search.searchEntityLocation lang toks
         (entities, err) <- case r of
           None t ->
@@ -608,9 +625,10 @@ fetchSymbolReferenceRanges
   -> Language
   -> [Text]
   -> Maybe Int
+  -> Maybe Text
   -> GleanBackend b
   -> IO ([LocationRange], Maybe ErrorLogger)
-fetchSymbolReferenceRanges env scsrepo lang toks limit b =
+fetchSymbolReferenceRanges env scsrepo lang toks limit filename b =
   backendRunHaxl b env $ do
     er <- symbolToAngleEntities lang toks
     let logDBs x = logError (gleanDBs b) <> x
@@ -621,7 +639,11 @@ fetchSymbolReferenceRanges env scsrepo lang toks limit b =
           withRepo entityRepo $ do
             let convert (targetFile, rspan) =
                   rangeSpanToLocationRange scsrepo targetFile rspan
-            uses <- searchWithLimit limit $ Query.findReferenceRangeSpan query
+            uses <- searchWithLimit limit $ case (filename, lang) of
+              -- Definition file filter only supported for Python
+              (Just filename, Language_Python) ->
+                Query.findReferenceRangeSpanByDef query filename
+              _ -> Query.findReferenceRangeSpan query
             mapM convert uses
         return (nubOrd $ concat ranges, fmap (logDBs . logError) searchErr)
 
@@ -687,9 +709,12 @@ searchRelated
 searchRelated env@Glass.Env{..} sym opts@RequestOptions{..}
     SearchRelatedRequest{..} =
   withSymbol "searchRelated" env opts sym $
-    \gleanDBs_ dbInfo (repo, lang, toks) ->
+    \gleanDBs_ dbInfo (repo, lang, toks, filters) ->
       let gleanDBs = filterDBs lang gleanDBs_ in
       backendRunHaxl GleanBackend{..} env $ do
+        case filters of
+          [] -> return ()
+          _ -> throwM (ServerException "searchRelated: filters not supported")
         entity <- searchFirstEntity lang toks
         withRepo (entityRepo entity) $ do
           (symPairs, desc, merr) <- case searchRelatedRequest_relatedBy of
@@ -922,8 +947,11 @@ searchRelatedNeighborhood
   -> IO RelatedNeighborhoodResult
 searchRelatedNeighborhood env@Glass.Env{..} sym opts@RequestOptions{..} req =
   withSymbol "searchRelatedNeighborhood" env opts sym $
-    \gleanDBs dbInfo (repo, lang, toks) ->
+    \gleanDBs dbInfo (repo, lang, toks, filters) ->
       backendRunHaxl GleanBackend{..} env $ do
+        case filters of
+          [] -> return ()
+          _ -> throwM (ServerException "searchRelatedNeighborhood: filters not supported")
         baseEntity <- searchFirstEntity lang toks
         let lang = entityLanguage (fst4 (decl baseEntity))
         (,Nothing) <$> searchNeighborhood limit req sym repo

glean/glass/Glean/Glass/Handler/Utils.hs

@@ -206,14 +206,14 @@ withSymbol
   -> SymbolId
   -> (  NonEmpty (GleanDBName, Glean.Repo)
      -> GleanDBInfo
-     -> (RepoName, Language, [Text])
+     -> (RepoName, Language, [Text], [SymbolFilter])
      -> IO (res, Maybe ErrorLogger))
   -> IO res
 withSymbol method env@Glass.Env{..} opts sym fn =
   fmap fst $ withRequest method env sym opts $ \dbInfo ->
     case symbolTokens sym of
       Left err -> throwM $ ServerException err
-      Right req@(repo, lang, _toks) -> do
+      Right req@(repo, lang, _toks, _filters) -> do
         dbs <- getGleanRepos tracer sourceControl repoMapping dbInfo repo
           (Just lang) (requestOptions_revision opts)
           (dbChooser repo opts) gleanDB

glean/glass/Glean/Glass/Query.hs

@@ -21,6 +21,7 @@ module Glean.Glass.Query
 
   -- * Finding references to declarations
   , findReferenceRangeSpan
+  , findReferenceRangeSpanByDef
 
   -- * Finding references for call hierarchy
   , findReferenceEntities
@@ -151,6 +152,43 @@ findReferenceRangeSpan ent =
       )
     ]
 
+-- | Entity-based find-references returning native range or bytespan filtered by
+-- the referenced symbols being defined in the given file.
+findReferenceRangeSpanByDef
+  :: Angle Code.Entity
+  -> Text
+  -> Angle (Src.File, Code.RangeSpan)
+findReferenceRangeSpanByDef ent filename =
+  vars $ \(reffile :: Angle Src.File)
+    (rangespan :: Angle Code.RangeSpan)
+    (deffile :: Angle Src.File)
+    (xref :: Angle Code.XRefLocation)
+    (targetLoc :: Angle Code.Location) ->
+      Angle.tuple (reffile, rangespan) `where_` [
+        wild .= predicate @Code.EntityReferences (
+          rec $
+              field @"target" ent $
+              field @"file" (asPredicate reffile) $
+              field @"range" rangespan
+            end
+          )
+        , deffile .= predicate @Src.File (string filename)
+        , wild .= predicate @Code.FileEntityXRefLocations (
+            rec $
+              field @"file" (asPredicate reffile) $
+              field @"xref" xref $
+              field @"entity" ent
+            end
+          )
+        , xref .= rec (
+            field @"target" targetLoc $
+            field @"source" rangespan
+          end)
+        , targetLoc .= rec (
+            field @"file" (asPredicate deffile)
+          end)
+      ]
+
 generatedEntityToIdlEntity
   :: Angle Code.Entity
   -> Angle (Code.Entity, Src.File)

glean/glass/Glean/Glass/SymbolId.hs

@@ -43,7 +43,7 @@ module Glean.Glass.SymbolId
   ,nativeSymbol) where
 
 import Control.Monad.Catch ( throwM, try )
-import Data.Maybe ( fromMaybe )
+import Data.Maybe ( fromMaybe, mapMaybe )
 import Data.Tuple ( swap )
 import Util.Text ( textShow )
 import Data.Text ( Text )
@@ -59,7 +59,8 @@ import Glean.Glass.Types as Glass
       Language(..),
       SymbolId(SymbolId),
       RepoName(..),
-      Name, )
+      Name,
+      SymbolFilter(..) )
 
 import Glean.Angle ( alt, Angle )
 import qualified Glean.Haxl.Repos as Glean
@@ -109,11 +110,12 @@ import Glean.Schema.CodeCsharp.Types as CSharp ( Entity(Entity_decl) )
 -- If we can't encode the entity, it is still useful to return the path and
 -- file, as this is enough to navigate with.
 --
--- Symbol IDs have 3 parts:
+-- Symbol IDs have 3 fixed parts and 1 optional part:
 --
 -- - scm repo (corpus in biggrep term), e.g. "fbsource"
 -- - language short code (e.g py or php or cpp)
 -- - entity encoding
+-- - [optional] filters as defined in the thrift schema
 --
 -- We uri encode the pieces and separate with / so they look like nice urls
 --
@@ -165,16 +167,32 @@ toSymbolQualifiedContainer entity = do
     Right (_, qualifiedName_container) -> Just qualifiedName_container
     Left _ -> Nothing
 
+parseFilter :: Text -> Maybe SymbolFilter
+parseFilter f = case Text.splitOn "=" f of
+  [k, v] -> case k of
+    "definition_file" -> Just (SymbolFilter_definition_file v)
+    _ -> Nothing
+  _ -> Nothing
+
+parseFilters :: Text -> [SymbolFilter]
+parseFilters = mapMaybe parseFilter . Text.split (=='&')
+
 -- | Tokenize a symbol (inverse of the intercalate "/")
 -- Leaves the path/qname/syms for further search.
-symbolTokens :: SymbolId -> Either Text (RepoName, Language, [Text])
+symbolTokens
+  :: SymbolId
+  -> Either Text (RepoName, Language, [Text], [SymbolFilter])
 symbolTokens (SymbolId symid)
   | (repo: code: pieces) <- tokens
   , Just lang <- fromShortCode code
-  = Right (RepoName repo, lang, map URI.decodeText pieces)
+  = Right (RepoName repo, lang, map URI.decodeText pieces, filters)
   | otherwise = Left $ "Invalid symbol: " <> symid
   where
-    tokens = Text.split (=='/') symid
+    (ident, filters) = case Text.split (=='?') symid of
+      [identOnly] -> (identOnly, [])
+      [identPart, filtersPart] -> (identPart, parseFilters filtersPart)
+      _ -> (symid, [])
+    tokens = Text.split (=='/') ident
 
 -- | SymbolID-encoded language, used for db name lookups
 shortCodeTable :: [(Language,Text)]

glean/glass/if/glass.thrift

@@ -187,6 +187,12 @@ struct DocumentSymbolsRequest {
 
 // response types
 
+// Filter for symbol based search results
+union SymbolFilter {
+  // Filters the symbol such that it is defined in the given file
+  1: string definition_file;
+}
+
 // Human-readable opaque, stable, globally unique symbol identifier
 typedef string SymbolId (hs.newtype)
 

glean/glass/test/regression/lib/Glean/Glass/Regression/Snapshot.hs

@@ -40,7 +40,6 @@ import qualified Data.ByteString.Lazy as B
 import qualified Data.Map.Strict as Map
 import qualified Data.Yaml as Yaml
 import Util.List ( uniq )
-import Data.Tuple.Extra
 import Data.Either.Extra
 
 import qualified Thrift.Protocol
@@ -226,8 +225,10 @@ validateCxxSymbols glassEnv req def = do
       refs = map Glass.referenceRangeSymbolX_sym
         (Glass.documentSymbolListXResult_references res)
       syms = uniq (defns ++ refs)
-      toks = map thd3 (filter ((== Language_Cpp) . snd3) -- only cxx ids thx
+      toks = map getTokens (filter ((== Language_Cpp) . getLanguage) -- only cxx ids thx
                 (rights (map Glass.symbolTokens syms)))
+      getLanguage (_, lang, _, _) = lang
+      getTokens (_, _, tokens, _) = tokens
   return (map Cxx.validateSymbolId toks)
 
 decodeObjectAsThriftJson