[Index] Introduce IndexSystem::foreachUnitTestSymbolReferencedByOutputPaths()

include/IndexStoreDB/Database/ReadTransaction.h

@@ -52,6 +52,9 @@ class INDEXSTOREDB_EXPORT ReadTransaction {
   /// Returns USR codes in batches.
   bool foreachUSROfGlobalSymbolKind(SymbolKind symKind, llvm::function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver);
 
+  /// Returns USR codes in batches.
+  bool foreachUSROfGlobalUnitTestSymbol(llvm::function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver);
+
   /// Returns USR codes in batches.
   bool findUSRsWithNameContaining(StringRef pattern,
                                   bool anchorStart, bool anchorEnd,

include/IndexStoreDB/Index/IndexSystem.h

@@ -124,6 +124,11 @@ class INDEXSTOREDB_EXPORT IndexSystem {
   bool foreachFileIncludedByFile(StringRef SourcePath,
                                               function_ref<bool(CanonicalFilePathRef TargetPath, unsigned Line)> Receiver);
 
+  /// Returns unit test class/method occurrences that are referenced from units associated with the provided output file paths.
+  /// \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<StringRef> FilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+
 private:
   IndexSystem(void *Impl) : Impl(Impl) {}
 

include/IndexStoreDB/Index/SymbolDataProvider.h

@@ -56,6 +56,9 @@ class SymbolDataProvider {
                                             SymbolRoleSet RoleSet,
                         function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) = 0;
 
+  virtual bool foreachUnitTestSymbolOccurrence(
+                        function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) = 0;
+
 private:
   virtual void anchor();
 };

include/IndexStoreDB/Index/SymbolIndex.h

@@ -83,6 +83,9 @@ class SymbolIndex {
   bool foreachCanonicalSymbolOccurrenceByKind(SymbolKind symKind, bool workspaceOnly,
                                               function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 
+  bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePath> FilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+
 private:
   void *Impl; // A SymbolIndexImpl.
 };

lib/Database/ReadTransaction.cpp

@@ -248,6 +248,14 @@ bool ReadTransaction::Implementation::foreachUSROfGlobalSymbolKind(SymbolKind sy
   return foreachUSROfGlobalSymbolKind(globalKindOpt.getValue(), receiver);
 }
 
+bool ReadTransaction::Implementation::foreachUSROfGlobalUnitTestSymbol(function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver) {
+  bool cont = foreachUSROfGlobalSymbolKind(GlobalSymbolKind::TestClassOrExtension, receiver);
+  if (cont) {
+    cont = foreachUSROfGlobalSymbolKind(GlobalSymbolKind::TestMethod, receiver);
+  }
+  return cont;
+}
+
 bool ReadTransaction::Implementation::foreachUSROfGlobalSymbolKind(GlobalSymbolKind globalKind,
                                                                    function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver) {
   auto &db = DBase->impl();
@@ -639,6 +647,10 @@ bool ReadTransaction::foreachUSROfGlobalSymbolKind(SymbolKind symKind, llvm::fun
   return Impl->foreachUSROfGlobalSymbolKind(symKind, std::move(receiver));
 }
 
+bool ReadTransaction::foreachUSROfGlobalUnitTestSymbol(llvm::function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver) {
+  return Impl->foreachUSROfGlobalUnitTestSymbol(std::move(receiver));
+}
+
 bool ReadTransaction::findUSRsWithNameContaining(StringRef pattern,
                                                  bool anchorStart, bool anchorEnd,
                                                  bool subsequence, bool ignoreCase,

lib/Database/ReadTransactionImpl.h

@@ -54,6 +54,7 @@ class ReadTransaction::Implementation {
   bool foreachProviderAndFileCodeReference(llvm::function_ref<bool(IDCode provider, IDCode pathCode, IDCode unitCode, llvm::sys::TimePoint<> modTime, IDCode moduleNameCode, bool isSystem)> receiver);
 
   bool foreachUSROfGlobalSymbolKind(SymbolKind symKind, llvm::function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver);
+  bool foreachUSROfGlobalUnitTestSymbol(llvm::function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver);
   bool foreachUSROfGlobalSymbolKind(GlobalSymbolKind globalSymKind, function_ref<bool(ArrayRef<IDCode> usrCodes)> receiver);
 
   bool findUSRsWithNameContaining(StringRef pattern,

lib/Index/IndexSystem.cpp

@@ -179,6 +179,9 @@ class IndexSystemImpl {
 
   bool foreachFileIncludedByFile(StringRef SourcePath,
                                  function_ref<bool(CanonicalFilePathRef TargetPath, unsigned Line)> Receiver);
+
+  bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<StringRef> FilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 };
 
 } // anonymous namespace
@@ -519,6 +522,15 @@ bool IndexSystemImpl::foreachFileIncludedByFile(StringRef SourcePath,
   return PathIndex->foreachFileIncludedByFile(canonSourcePath, Receiver);
 }
 
+bool IndexSystemImpl::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<StringRef> FilePaths, function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  SmallVector<CanonicalFilePath, 8> canonPaths;
+  canonPaths.reserve(FilePaths.size());
+  for (StringRef path : FilePaths) {
+    canonPaths.push_back(PathIndex->getCanonicalPath(path));
+  }
+  return SymIndex->foreachUnitTestSymbolReferencedByOutputPaths(canonPaths, std::move(Receiver));
+}
+
 //===----------------------------------------------------------------------===//
 // IndexSystem
 //===----------------------------------------------------------------------===//
@@ -696,3 +708,8 @@ bool IndexSystem::foreachFileIncludedByFile(StringRef SourcePath,
                                                 function_ref<bool(CanonicalFilePathRef TargetPath, unsigned Line)> Receiver) {
   return IMPL->foreachFileIncludedByFile(SourcePath, Receiver);
 }
+
+bool IndexSystem::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<StringRef> FilePaths,
+    function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  return IMPL->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));
+}

lib/Index/StoreSymbolRecord.cpp

@@ -376,3 +376,29 @@ bool StoreSymbolRecord::foreachRelatedSymbolOccurrenceByUSR(ArrayRef<db::IDCode>
 
   return !Err && Finished;
 }
+
+bool StoreSymbolRecord::foreachUnitTestSymbolOccurrence(function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  bool Finished = true;
+  bool Err = doForData([&](IndexRecordReader &Reader) {
+    SmallVector<IndexRecordSymbol, 8> FoundDecls;
+    auto filter = [&](IndexRecordSymbol recSym, bool &stop) -> bool {
+      auto symInfo = getSymbolInfo(recSym);
+      return symInfo.Properties.contains(SymbolProperty::UnitTest);
+    };
+    auto receiver = [&](IndexRecordSymbol sym) {
+      FoundDecls.push_back(sym);
+    };
+    Reader.searchSymbols(filter, receiver);
+    if (FoundDecls.empty())
+      return;
+
+    // Return all occurrences.
+    auto Pred = [](IndexRecordOccurrence) -> bool { return true; };
+    PredOccurrenceConverter Converter(*this, Pred, Receiver);
+    Finished = Reader.foreachOccurrence(/*symbolsFilter=*/FoundDecls,
+                                        /*relatedSymbolsFilter=*/None,
+                                        Converter);
+  });
+
+  return !Err && Finished;
+}

lib/Index/StoreSymbolRecord.h

@@ -82,6 +82,10 @@ class StoreSymbolRecord : public SymbolDataProvider {
   virtual bool foreachRelatedSymbolOccurrenceByUSR(ArrayRef<db::IDCode> USRs,
                                             SymbolRoleSet RoleSet,
                function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) override;
+
+  virtual bool foreachUnitTestSymbolOccurrence(
+               function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) override;
+
 };
 
 } // namespace index

lib/Index/SymbolIndex.cpp

@@ -78,6 +78,8 @@ class SymbolIndexImpl {
   size_t countOfCanonicalSymbolsWithKind(SymbolKind symKind, bool workspaceOnly);
   bool foreachCanonicalSymbolOccurrenceByKind(SymbolKind symKind, bool workspaceOnly,
                                               function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+  bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePath> FilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 
 private:
   bool foreachCanonicalSymbolImpl(bool workspaceOnly,
@@ -88,7 +90,8 @@ class SymbolIndexImpl {
                                             function_ref<bool(SymbolOccurrenceRef)> Receiver);
   std::vector<SymbolDataProviderRef> lookupProvidersForUSR(StringRef USR, SymbolRoleSet roles, SymbolRoleSet relatedRoles);
   std::vector<std::pair<SymbolDataProviderRef, bool>> findCanonicalProvidersForUSR(IDCode usrCode);
-  SymbolDataProviderRef createProviderForCode(IDCode providerCode, ReadTransaction &reader);
+  SymbolDataProviderRef createVisibleProviderForCode(IDCode providerCode, ReadTransaction &reader);
+  SymbolDataProviderRef createProviderForCode(IDCode providerCode, ReadTransaction &reader, function_ref<bool(ReadTransaction &, const UnitInfo &)> unitFilter);
 };
 
 } // anonymous namespace
@@ -162,7 +165,13 @@ void SymbolIndexImpl::dumpProviderFileAssociations(raw_ostream &OS) {
   });
 }
 
-SymbolDataProviderRef SymbolIndexImpl::createProviderForCode(IDCode providerCode, ReadTransaction &reader) {
+SymbolDataProviderRef SymbolIndexImpl::createVisibleProviderForCode(IDCode providerCode, ReadTransaction &reader) {
+  return createProviderForCode(providerCode, reader, [&](ReadTransaction &reader, const UnitInfo &unitInfo) -> bool {
+    return VisibilityChecker->isUnitVisible(unitInfo, reader);
+  });
+}
+
+SymbolDataProviderRef SymbolIndexImpl::createProviderForCode(IDCode providerCode, ReadTransaction &reader, function_ref<bool(ReadTransaction &, const UnitInfo &)> unitFilter) {
   StringRef recordName = reader.getProviderName(providerCode);
   if (recordName.empty()) {
     ++NumMissingProvidersLookedUp;
@@ -175,7 +184,7 @@ SymbolDataProviderRef SymbolIndexImpl::createProviderForCode(IDCode providerCode
     auto unitInfo = reader.getUnitInfo(unitCode);
     if (unitInfo.isInvalid())
       return true;
-    if (!VisibilityChecker->isUnitVisible(unitInfo, reader))
+    if (!unitFilter(reader, unitInfo))
       return true;
 
     if (!providerKind.hasValue()) {
@@ -206,7 +215,7 @@ SymbolIndexImpl::lookupProvidersForUSR(StringRef USR, SymbolRoleSet roles, Symbo
   std::vector<SymbolDataProviderRef> providers;
   ReadTransaction reader(DBase);
   reader.lookupProvidersForUSR(USR, roles, relatedRoles, [&](IDCode providerCode, SymbolRoleSet roles, SymbolRoleSet relatedRoles) -> bool {
-    if (auto prov = createProviderForCode(providerCode, reader))
+    if (auto prov = createVisibleProviderForCode(providerCode, reader))
       providers.push_back(prov);
     return true;
   });
@@ -279,7 +288,7 @@ bool SymbolIndexImpl::foreachCanonicalSymbolImpl(bool workspaceOnly,
           if (provInfo.IsInvisible)
             return provInfo;
           if (!provInfo.Provider) {
-            provInfo.Provider = createProviderForCode(provCode, reader);
+            provInfo.Provider = createVisibleProviderForCode(provCode, reader);
             if (!provInfo.Provider)
               provInfo.IsInvisible = true;
           }
@@ -442,13 +451,54 @@ SymbolIndexImpl::findCanonicalProvidersForUSR(IDCode usrCode) {
     bool isCanon = provCodeAndHasDef.second;
     if (!isCanon && foundCanon)
       break;
-    if (auto prov = createProviderForCode(provCode, reader))
+    if (auto prov = createVisibleProviderForCode(provCode, reader))
       foundProvs.emplace_back(std::move(prov), isCanon);
     foundCanon |= isCanon;
   }
   return foundProvs;
 }
 
+bool SymbolIndexImpl::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePath> outFilePaths, function_ref<bool(SymbolOccurrenceRef Occur)> receiver) {
+  std::vector<SymbolDataProviderRef> providers;
+  {
+    ReadTransaction reader(DBase);
+
+    std::unordered_set<IDCode> providerCodes;
+    reader.foreachUSROfGlobalUnitTestSymbol([&](ArrayRef<IDCode> usrCodes) -> bool {
+      for (IDCode usrCode : usrCodes) {
+        reader.lookupProvidersForUSR(usrCode, None, None, [&](IDCode providerCode, SymbolRoleSet roles, SymbolRoleSet relatedRoles) -> bool {
+          providerCodes.insert(providerCode);
+          return true;
+        });
+      }
+      return true;
+    });
+
+    if (providerCodes.empty()) {
+      return true;
+    }
+
+    std::unordered_set<IDCode> outFileCodes;
+    for (const CanonicalFilePath &path : outFilePaths) {
+      outFileCodes.insert(reader.getFilePathCode(path));
+    }
+    for (IDCode providerCode : providerCodes) {
+      auto provider = createProviderForCode(providerCode, reader, [&](ReadTransaction &reader, const UnitInfo &unitInfo) -> bool {
+        return outFileCodes.count(unitInfo.OutFileCode);
+      });
+      if (provider) {
+        providers.push_back(std::move(provider));
+      }
+    }
+  }
+
+  for (SymbolDataProviderRef provider : providers) {
+    bool cont = provider->foreachUnitTestSymbolOccurrence(receiver);
+    if (!cont) return false;
+  }
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // SymbolIndex
 //===----------------------------------------------------------------------===//
@@ -527,3 +577,8 @@ bool SymbolIndex::foreachCanonicalSymbolOccurrenceByKind(SymbolKind symKind, boo
                                                          function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
   return IMPL->foreachCanonicalSymbolOccurrenceByKind(symKind, workspaceOnly, std::move(Receiver));
 }
+
+bool SymbolIndex::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePath> FilePaths,
+    function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  return IMPL->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));
+}