Add functions to retrieve all tests within a workspace or within a given source file

Sources/IndexStoreDB/IndexStoreDB.swift

@@ -381,6 +381,28 @@ public final class IndexStoreDB {
       return body(Symbol(symbol))
     }
   }
+
+  /// Returns all unit test symbol in unit files that reference one of the main files in `mainFilePaths`.
+  public func unitTests(referencedByMainFiles mainFilePaths: [String]) -> [SymbolOccurrence] {
+    var result: [SymbolOccurrence] = []
+    let cMainFiles: [UnsafePointer<CChar>] = mainFilePaths.map { UnsafePointer($0.withCString(strdup)!) }
+    defer { for cPath in cMainFiles { free(UnsafeMutablePointer(mutating: cPath)) } }
+    indexstoredb_index_unit_tests_referenced_by_main_files(impl, cMainFiles, cMainFiles.count) { symbol in
+      result.append(SymbolOccurrence(symbol))
+      return true
+    }
+    return result
+  }
+
+  /// Returns all unit test symbols in the index.
+  public func unitTests() -> [SymbolOccurrence] {
+    var result: [SymbolOccurrence] = []
+    indexstoredb_index_unit_tests(impl) { symbol in
+      result.append(SymbolOccurrence(symbol))
+      return true
+    }
+    return result
+  }
 }
 
 public protocol IndexStoreLibraryProvider {

include/CIndexStoreDB/CIndexStoreDB.h

@@ -522,6 +522,36 @@ indexstoredb_index_includes_of_unit(
   const char *_Nonnull unitName,
   _Nonnull indexstoredb_unit_includes_receiver receiver);
 
+/// Calls `receiver` for every unit test symbol in unit files that reference
+/// one of the main files in `mainFilePaths`.
+///
+/// \param index An IndexStoreDB object which contains the symbols.
+/// \param mainFilePaths File paths to search for unit tests
+/// \param count Number of elements in `mainFilePaths`.
+/// \param receiver A function to be called for each unit tests. If the receiver
+/// returns `false`, iteration is stopped and the function returns `true`.
+/// \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+INDEXSTOREDB_PUBLIC bool
+indexstoredb_index_unit_tests_referenced_by_main_files(
+  _Nonnull indexstoredb_index_t index,
+   const char *_Nonnull const *_Nonnull mainFilePaths,
+   size_t count,
+  _Nonnull indexstoredb_symbol_occurrence_receiver_t receiver
+);
+
+/// Calls `receiver` for every unit test symbol in the index.
+///
+/// \param index An IndexStoreDB object which contains the symbols.
+/// \param receiver A function to be called for each unit tests. If the receiver
+/// returns `false`, iteration is stopped and the function returns `true`.
+/// \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+INDEXSTOREDB_PUBLIC bool
+indexstoredb_index_unit_tests(
+  _Nonnull indexstoredb_index_t index,
+  _Nonnull indexstoredb_symbol_occurrence_receiver_t receiver
+);
+
+
 INDEXSTOREDB_END_DECLS
 
 #endif

include/IndexStoreDB/Index/IndexSystem.h

@@ -162,6 +162,21 @@ class INDEXSTOREDB_EXPORT IndexSystem {
   bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths,
       function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 
+
+  /// Calls `receiver` for every unit test symbol in unit files that reference
+  /// one of the main files in `mainFilePaths`.
+  ///
+  ///  \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbolReferencedByMainFiles(
+      ArrayRef<StringRef> mainFilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> receiver
+  );
+
+  /// Calls `receiver` for every unit test symbol in the index.
+  ///
+  ///  \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
+
 private:
   IndexSystem(void *Impl) : Impl(Impl) {}
 

include/IndexStoreDB/Index/SymbolIndex.h

@@ -89,6 +89,18 @@ class SymbolIndex {
   bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths,
       function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 
+  /// Calls `receiver` for every unit test symbol in unit files that reference
+  /// one of the main files in `mainFilePaths`.
+  ///
+  ///  \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbolReferencedByMainFiles(
+      ArrayRef<CanonicalFilePath> mainFilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> receiver
+  );
+  /// Calls `receiver` for every unit test symbol in the index.
+  ///
+  ///  \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
 private:
   void *Impl; // A SymbolIndexImpl.
 };

lib/CIndexStoreDB/CIndexStoreDB.cpp

@@ -15,6 +15,7 @@
 #include "IndexStoreDB/Index/IndexStoreLibraryProvider.h"
 #include "IndexStoreDB/Index/IndexSystem.h"
 #include "IndexStoreDB/Index/IndexSystemDelegate.h"
+#include "IndexStoreDB/Support/Path.h"
 #include "IndexStoreDB/Core/Symbol.h"
 #include "indexstore/IndexStoreCXX.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
@@ -598,4 +599,32 @@ indexstoredb_index_includes_of_unit(
   });
 }
 
+bool
+indexstoredb_index_unit_tests_referenced_by_main_files(
+  _Nonnull indexstoredb_index_t index,
+   const char *_Nonnull const *_Nonnull cMainFilePaths,
+   size_t count,
+  _Nonnull indexstoredb_symbol_occurrence_receiver_t receiver
+) {
+  auto obj = (Object<std::shared_ptr<IndexSystem>> *)index;
+  SmallVector<StringRef, 2> mainFilePaths;
+  for (size_t i = 0; i < count; ++i) {
+    mainFilePaths.push_back(cMainFilePaths[i]);
+  }
+  return obj->value->foreachUnitTestSymbolReferencedByMainFiles(mainFilePaths, [&](SymbolOccurrenceRef occur) -> bool {
+    return receiver((indexstoredb_symbol_occurrence_t)occur.get());
+  });
+}
+
+bool
+indexstoredb_index_unit_tests(
+  _Nonnull indexstoredb_index_t index,
+  _Nonnull indexstoredb_symbol_occurrence_receiver_t receiver
+) {
+  auto obj = (Object<std::shared_ptr<IndexSystem>> *)index;
+  return obj->value->foreachUnitTestSymbol([&](SymbolOccurrenceRef occur) -> bool {
+    return receiver((indexstoredb_symbol_occurrence_t)occur.get());
+  });
+}
+
 ObjectBase::~ObjectBase() {}

lib/Index/IndexSystem.cpp

@@ -219,6 +219,20 @@ class IndexSystemImpl {
 
   bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths,
       function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+
+  /// Calls `receiver` for every unit test symbol in unit files that reference
+  /// one of the main files in `mainFilePaths`.
+  ///
+  ///  \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbolReferencedByMainFiles(
+      ArrayRef<StringRef> mainFilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> receiver
+  );
+
+  /// Calls `receiver` for every unit test symbol in the index.
+  ///
+  ///  \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
 };
 
 } // anonymous namespace
@@ -598,6 +612,21 @@ bool IndexSystemImpl::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<Cano
   return SymIndex->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));
 }
 
+bool IndexSystemImpl::foreachUnitTestSymbolReferencedByMainFiles(
+    ArrayRef<StringRef> mainFilePaths,
+    function_ref<bool(SymbolOccurrenceRef Occur)> receiver
+) {
+  std::vector<CanonicalFilePath> canonicalMainFilesPaths;
+  for (StringRef mainFilePath : mainFilePaths) {
+    canonicalMainFilesPaths.push_back(PathIndex->getCanonicalPath(mainFilePath));
+  }
+  return SymIndex->foreachUnitTestSymbolReferencedByMainFiles(canonicalMainFilesPaths, std::move(receiver));
+}
+
+bool IndexSystemImpl::foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver) {
+  return SymIndex->foreachUnitTestSymbol(std::move(receiver));
+}
+
 //===----------------------------------------------------------------------===//
 // IndexSystem
 //===----------------------------------------------------------------------===//
@@ -804,3 +833,14 @@ bool IndexSystem::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<Canonica
     function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
   return IMPL->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));
 }
+
+bool IndexSystem::foreachUnitTestSymbolReferencedByMainFiles(
+   ArrayRef<StringRef> mainFilePaths,
+   function_ref<bool(SymbolOccurrenceRef Occur)> receiver
+) {
+  return IMPL->foreachUnitTestSymbolReferencedByMainFiles(mainFilePaths, std::move(receiver));
+}
+
+bool IndexSystem::foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver) {
+  return IMPL->foreachUnitTestSymbol(std::move(receiver));
+}

lib/Index/SymbolIndex.cpp

@@ -86,7 +86,28 @@ class SymbolIndexImpl {
   bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths,
       function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 
+  /// Calls `receiver` for every unit test symbol in unit files that reference
+  /// one of the main files in `mainFilePaths`.
+  ///
+  /// \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbolReferencedByMainFiles(
+      ArrayRef<CanonicalFilePath> mainFilePaths,
+      function_ref<bool(SymbolOccurrenceRef Occur)> receiver
+  );
+  /// Calls `receiver` for every unit test symbol in the index.
+  ///
+  /// \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
+
 private:
+  /// Returns all the providers in the index that contain test cases and satisfy `unitFilter`.
+  std::vector<SymbolDataProviderRef> providersContainingTestCases(ReadTransaction &reader, function_ref<bool(const UnitInfo &)> unitFilter);
+
+  /// Calls `receiver` for every unit test contained by a provider in `providers`.
+  ///
+  /// \returns `false` if the receiver returned `false` to stop receiving symbols, `true` otherwise.
+  bool foreachUnitTestSymbolOccurrence(const std::vector<SymbolDataProviderRef> &providers, function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
+
   bool foreachCanonicalSymbolImpl(bool workspaceOnly,
                                   function_ref<bool(ReadTransaction &, function_ref<bool(ArrayRef<IDCode> usrCode)> usrConsumer)> usrProducer,
                                   function_ref<bool(SymbolDataProviderRef, std::vector<std::pair<IDCode, bool>> USRs)> receiver);
@@ -518,30 +539,57 @@ bool SymbolIndexImpl::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<Cano
   {
     ReadTransaction reader(DBase);
 
-    SmallVector<IDCode, 16> providerCodes;
-    reader.foreachProviderContainingTestSymbols([&](IDCode providerCode) -> bool {
-      providerCodes.push_back(providerCode);
-      return true;
-    });
-
-    if (providerCodes.empty()) {
-      return true;
-    }
-
     std::unordered_set<IDCode> outFileCodes;
     for (const CanonicalFilePathRef &path : outFilePaths) {
       outFileCodes.insert(reader.getFilePathCode(path));
     }
-    for (IDCode providerCode : providerCodes) {
-      auto provider = createProviderForCode(providerCode, reader, [&](const UnitInfo &unitInfo) -> bool {
-        return outFileCodes.count(unitInfo.OutFileCode);
-      });
-      if (provider) {
-        providers.push_back(std::move(provider));
-      }
+
+    providers = providersContainingTestCases(reader, [&](const UnitInfo &unitInfo) -> bool {
+      return outFileCodes.count(unitInfo.OutFileCode);
+    });
+  }
+
+  return foreachUnitTestSymbolOccurrence(providers, receiver);
+}
+
+bool SymbolIndexImpl::foreachUnitTestSymbolReferencedByMainFiles(ArrayRef<CanonicalFilePath> mainFilePaths, function_ref<bool(SymbolOccurrenceRef Occur)> receiver) {
+  std::vector<SymbolDataProviderRef> providers;
+  {
+    ReadTransaction reader(DBase);
+
+    std::unordered_set<IDCode> fileCodes;
+    for (const CanonicalFilePathRef &path : mainFilePaths) {
+      fileCodes.insert(reader.getFilePathCode(path));
     }
+
+    providers = providersContainingTestCases(reader, [&](const UnitInfo &unitInfo) -> bool {
+      return fileCodes.count(unitInfo.MainFileCode);
+    });
+  }
+  return foreachUnitTestSymbolOccurrence(providers, receiver);
+}
+
+bool SymbolIndexImpl::foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver) {
+  std::vector<SymbolDataProviderRef> providers;
+  {
+    ReadTransaction reader(DBase);
+    providers = providersContainingTestCases(reader, [&](const UnitInfo &unitInfo) -> bool { return true; });
   }
+  return foreachUnitTestSymbolOccurrence(providers, receiver);
+}
+
+std::vector<SymbolDataProviderRef> SymbolIndexImpl::providersContainingTestCases(ReadTransaction &reader, function_ref<bool(const UnitInfo &)> unitFilter) {
+  std::vector<SymbolDataProviderRef> providers;
+  reader.foreachProviderContainingTestSymbols([&](IDCode providerCode) -> bool {
+    if (auto provider = createProviderForCode(providerCode, reader, unitFilter)) {
+      providers.push_back(std::move(provider));
+    }
+    return true;
+  });
+  return providers;
+}
 
+bool SymbolIndexImpl::foreachUnitTestSymbolOccurrence(const std::vector<SymbolDataProviderRef> &providers, function_ref<bool(SymbolOccurrenceRef Occur)> receiver) {
   for (SymbolDataProviderRef provider : providers) {
     bool cont = provider->foreachUnitTestSymbolOccurrence(receiver);
     if (!cont) return false;
@@ -637,3 +685,14 @@ bool SymbolIndex::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<Canonica
     function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
   return IMPL->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));
 }
+
+bool SymbolIndex::foreachUnitTestSymbolReferencedByMainFiles(
+     ArrayRef<CanonicalFilePath> mainFilePaths,
+     function_ref<bool(SymbolOccurrenceRef Occur)> receiver
+ ) {
+  return IMPL->foreachUnitTestSymbolReferencedByMainFiles(mainFilePaths, std::move(receiver));
+}
+
+bool SymbolIndex::foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver) {
+  return IMPL->foreachUnitTestSymbol(std::move(receiver));
+}