[Index] Add APIs to provide the recorded #includes of a unit file

Sources/IndexStoreDB/IndexStoreDB.swift

@@ -195,8 +195,8 @@ public final class IndexStoreDB {
   @discardableResult
   public func forEachMainFileContainingFile(path: String, crossLanguage: Bool, body: @escaping (String) -> Bool) -> Bool {
     let fromSwift = path.hasSuffix(".swift")
-    return indexstoredb_index_main_files_containing_file(impl, path) { mainFile in
-      let mainFileStr = String(cString: mainFile)
+    return indexstoredb_index_units_containing_file(impl, path) { unit in
+      let mainFileStr = String(cString: indexstoredb_unit_info_main_file_path(unit))
       let toSwift = mainFileStr.hasSuffix(".swift")
       if !crossLanguage && fromSwift != toSwift {
         return true // continue
@@ -214,6 +214,59 @@ public final class IndexStoreDB {
     return result
   }
 
+  @discardableResult
+  public func forEachUnitNameContainingFile(path: String, body: @escaping (String) -> Bool) -> Bool {
+    return indexstoredb_index_units_containing_file(impl, path) { unit in
+      let unitName = String(cString: indexstoredb_unit_info_unit_name(unit))
+      return body(unitName)
+    }
+  }
+
+  public func unitNamesContainingFile(path: String) -> [String] {
+    var result: [String] = []
+    forEachUnitNameContainingFile(path: path) { unitName in
+      result.append(unitName)
+      return true
+    }
+    return result
+  }
+
+  /// A recorded header `#include` from a unit file.
+  public struct UnitIncludeEntry: Equatable {
+    /// The path where the `#include` was added.
+    public let sourcePath: String
+    /// The path that the `#include` resolved to.
+    public let targetPath: String
+    /// the line where the `#include` was added.
+    public let line: Int
+
+    public init(sourcePath: String, targetPath: String, line: Int) {
+      self.sourcePath = sourcePath
+      self.targetPath = targetPath
+      self.line = line
+    }
+  }
+
+  /// Iterates over recorded `#include`s of a unit.
+  @discardableResult
+  public func forEachIncludeOfUnit(unitName: String, body: @escaping (UnitIncludeEntry) -> Bool) -> Bool {
+    return indexstoredb_index_includes_of_unit(impl, unitName) { sourcePath, targetPath, line in
+      let sourcePathStr = String(cString: sourcePath)
+      let targetPathStr = String(cString: targetPath)
+      return body(UnitIncludeEntry(sourcePath: sourcePathStr, targetPath: targetPathStr, line: line))
+    }
+  }
+
+  /// Returns the recorded `#include`s of a unit.
+  public func includesOfUnit(unitName: String) -> [UnitIncludeEntry] {
+    var result: [UnitIncludeEntry] = []
+    forEachIncludeOfUnit(unitName: unitName) { entry in
+      result.append(entry)
+      return true
+    }
+    return result
+  }
+
   /// Iterates over the name of every symbol in the index.
   ///
   /// - Parameter body: A closure to be called for each symbol. The closure should return true to

Tests/INPUTS/MainFiles/main1.c

@@ -1,5 +1,5 @@
-#include "shared.h"
-#include "uniq1.h"
+#include "shared.h" /*include_main1_shared*/
+#include "uniq1.h" /*include_main1_uniq1*/
 
 void main1(void) {
   /*main1*/

Tests/IndexStoreDBTests/IndexTests.swift

@@ -280,6 +280,26 @@ final class IndexTests: XCTestCase {
     XCTAssertEqual(mainFiles(unknown, false), [])
   }
 
+  func testUnitIncludes() throws {
+    guard let ws = try staticTibsTestWorkspace(name: "MainFiles") else { return }
+    try ws.buildAndIndex()
+    let index = ws.index
+
+    let main1 = ws.testLoc("main1").url.path
+    let uniq1 = ws.testLoc("uniq1").url.path
+    let shared = ws.testLoc("shared").url.path
+
+    let units = index.unitNamesContainingFile(path: main1)
+    XCTAssertEqual(units.count, 1)
+    let main1Unit = try XCTUnwrap(units.first)
+
+    let includes = index.includesOfUnit(unitName: main1Unit)
+    XCTAssertEqual(includes, [
+      IndexStoreDB.UnitIncludeEntry(sourcePath: main1, targetPath: shared, line: ws.testLoc("include_main1_shared").line),
+      IndexStoreDB.UnitIncludeEntry(sourcePath: main1, targetPath: uniq1, line: ws.testLoc("include_main1_uniq1").line),
+    ])
+  }
+
   func testAllSymbolNames() throws {
     guard let ws = try staticTibsTestWorkspace(name: "proj1") else { return }
     try ws.buildAndIndex()

Tests/IndexStoreDBTests/XCTestManifests.swift

@@ -24,6 +24,7 @@ extension IndexTests {
         ("testMainFilesContainingFile", testMainFilesContainingFile),
         ("testMixedLangTarget", testMixedLangTarget),
         ("testSwiftModules", testSwiftModules),
+        ("testUnitIncludes", testUnitIncludes),
         ("testWaitUntilDoneInitializing", testWaitUntilDoneInitializing),
     ]
 }

include/CIndexStoreDB/CIndexStoreDB.h

@@ -57,6 +57,7 @@ typedef void *indexstoredb_symbol_occurrence_t;
 typedef void *indexstoredb_error_t;
 typedef void *indexstoredb_symbol_location_t;
 typedef void *indexstoredb_symbol_relation_t;
+typedef void *indexstoredb_unit_info_t;
 
 typedef enum {
   INDEXSTOREDB_SYMBOL_ROLE_DECLARATION = 1 << 0,
@@ -135,7 +136,10 @@ typedef bool(^indexstoredb_symbol_name_receiver)(const char *_Nonnull);
 typedef void(^indexstoredb_delegate_event_receiver_t)(_Nonnull indexstoredb_delegate_event_t);
 
 /// Returns true to continue.
-typedef bool(^indexstoredb_path_receiver)(const char *_Nonnull);
+typedef bool(^indexstoredb_unit_info_receiver)(_Nonnull indexstoredb_unit_info_t);
+
+/// Returns true to continue.
+typedef bool(^indexstoredb_unit_includes_receiver)(const char *_Nonnull sourcePath, const char *_Nonnull targetPath, size_t line);
 
 /// Creates an index for the given raw index data in \p storePath.
 ///
@@ -352,20 +356,41 @@ indexstoredb_symbol_occurrence_relations(_Nonnull indexstoredb_symbol_occurrence
 INDEXSTOREDB_PUBLIC indexstoredb_symbol_kind_t
 indexstoredb_symbol_kind(_Nonnull indexstoredb_symbol_t);
 
-/// Iterates over the compilation units that contain \p path and return their main file.
+/// Returns the main file path of a unit info object.
+///
+/// The main file is typically the one that e.g. a build system would have explicit knowledge of.
+INDEXSTOREDB_PUBLIC const char *_Nonnull
+indexstoredb_unit_info_main_file_path(_Nonnull indexstoredb_unit_info_t);
+
+/// Returns the unit name of a unit info object.
+INDEXSTOREDB_PUBLIC const char *_Nonnull
+indexstoredb_unit_info_unit_name(_Nonnull indexstoredb_unit_info_t);
+
+/// Iterates over the compilation units that contain \p path and return their units.
 ///
-/// This can be used to find the main files that include a given header. The main file is typically
-/// the one that e.g. a build system would have explicit knowledge of.
+/// This can be used to find information for units that include a given header.
 ///
 /// \param index An IndexStoreDB object which contains the symbols.
 /// \param path The source file to search for.
-/// \param receiver A function to be called for each main file path. The string pointer is only valid for
+/// \param receiver A function to be called for each unit. The pointer is only valid for
 /// the duration of the call. The function should return a true to continue iterating.
 INDEXSTOREDB_PUBLIC bool
-indexstoredb_index_main_files_containing_file(
+indexstoredb_index_units_containing_file(
   _Nonnull indexstoredb_index_t index,
   const char *_Nonnull path,
-  _Nonnull indexstoredb_path_receiver receiver);
+  _Nonnull indexstoredb_unit_info_receiver receiver);
+
+/// Iterates over recorded `#include`s of a unit.
+///
+/// \param index An IndexStoreDB object which contains the symbols.
+/// \param unitName The unit name to search for.
+/// \param receiver A function to be called for each include entry. The pointers are only valid for
+/// the duration of the call. The function should return a true to continue iterating.
+INDEXSTOREDB_PUBLIC bool
+indexstoredb_index_includes_of_unit(
+  _Nonnull indexstoredb_index_t index,
+  const char *_Nonnull unitName,
+  _Nonnull indexstoredb_unit_includes_receiver receiver);
 
 INDEXSTOREDB_END_DECLS
 

include/IndexStoreDB/Index/FilePathIndex.h

@@ -72,6 +72,9 @@ class FilePathIndex {
   bool foreachFileIncludedByFile(CanonicalFilePathRef sourcePath,
                             function_ref<bool(CanonicalFilePathRef TargetPath, unsigned Line)> Receiver);
 
+  bool foreachIncludeOfUnit(StringRef unitName,
+                            function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver);
+
 private:
   void *Impl; // A FileIndexImpl.
 };

include/IndexStoreDB/Index/IndexSystem.h

@@ -139,6 +139,9 @@ class INDEXSTOREDB_EXPORT IndexSystem {
   bool foreachFileIncludedByFile(StringRef SourcePath,
                                               function_ref<bool(CanonicalFilePathRef TargetPath, unsigned Line)> Receiver);
 
+  bool foreachIncludeOfUnit(StringRef unitName,
+                            function_ref<bool(CanonicalFilePathRef sourcePath, 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<CanonicalFilePathRef> FilePaths,

include/IndexStoreDB/Support/Path.h

@@ -22,13 +22,13 @@ namespace IndexStoreDB {
   class CanonicalFilePathRef;
 
 class CanonicalFilePath {
-  SmallString<128> Path;
+  std::string Path;
 
 public:
   CanonicalFilePath() = default;
   inline CanonicalFilePath(CanonicalFilePathRef CanonPath);
 
-  StringRef getPath() const { return Path; }
+  const std::string &getPath() const { return Path; }
   bool empty() const { return Path.empty(); }
 
   friend bool operator==(CanonicalFilePath LHS, CanonicalFilePath RHS) {

lib/CIndexStoreDB/CIndexStoreDB.cpp

@@ -422,14 +422,38 @@ static indexstoredb_symbol_kind_t toCSymbolKind(SymbolKind K) {
   }
 }
 
+const char *
+indexstoredb_unit_info_main_file_path(indexstoredb_unit_info_t info) {
+  auto obj = (const StoreUnitInfo *)info;
+  return obj->MainFilePath.getPath().c_str();
+}
+
+const char *
+indexstoredb_unit_info_unit_name(indexstoredb_unit_info_t info) {
+  auto obj = (const StoreUnitInfo *)info;
+  return obj->UnitName.c_str();
+}
+
 bool
-indexstoredb_index_main_files_containing_file(
+indexstoredb_index_units_containing_file(
   indexstoredb_index_t index,
   const char *path,
-  indexstoredb_path_receiver receiver)
+  indexstoredb_unit_info_receiver receiver)
 {
   auto obj = (IndexStoreDBObject<std::shared_ptr<IndexSystem>> *)index;
   return obj->value->foreachMainUnitContainingFile(path, [&](const StoreUnitInfo &unitInfo) -> bool {
-    return receiver(unitInfo.MainFilePath.getPath().str().c_str());
+    return receiver((indexstoredb_unit_info_receiver)&unitInfo);
+  });
+}
+
+bool
+indexstoredb_index_includes_of_unit(
+  indexstoredb_index_t index,
+  const char *unitName,
+  indexstoredb_unit_includes_receiver receiver)
+{
+  auto obj = (IndexStoreDBObject<std::shared_ptr<IndexSystem>> *)index;
+  return obj->value->foreachIncludeOfUnit(unitName, [&](CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)->bool {
+    return receiver(sourcePath.getPath().str().c_str(), targetPath.getPath().str().c_str(), line);
   });
 }

lib/Index/FilePathIndex.cpp

@@ -76,6 +76,9 @@ class FileIndexImpl {
   bool foreachFileIncludedByFile(CanonicalFilePathRef sourcePath,
                                  function_ref<bool(CanonicalFilePathRef targetPath, unsigned line)> Receiver);
 
+  bool foreachIncludeOfUnit(StringRef unitName,
+                            function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver);
+
   bool foreachIncludeOfStoreUnitContainingFile(CanonicalFilePathRef filePath,
                function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver);
 };
@@ -223,6 +226,23 @@ bool FileIndexImpl::foreachFileIncludedByFile(CanonicalFilePathRef inputSourcePa
   });
 }
 
+bool FileIndexImpl::foreachIncludeOfUnit(StringRef unitName,
+        function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver) {
+  std::string error;
+  IndexUnitReader storeUnit(*IdxStore, unitName, error);
+  if (!storeUnit.isValid())
+    return true;
+  StringRef workDir = storeUnit.getWorkingDirectory();
+  return storeUnit.foreachInclude([&](IndexUnitInclude inc) -> bool {
+    StringRef sourcePath = inc.getSourcePath();
+    StringRef targetPath = inc.getTargetPath();
+    unsigned line = inc.getSourceLine();
+    CanonicalFilePath fullSourcePath = getCanonicalPath(sourcePath, workDir);
+    CanonicalFilePath fullTargetPath = getCanonicalPath(targetPath, workDir);
+    return receiver(fullSourcePath, fullTargetPath, line);
+  });
+}
+
 bool FileIndexImpl::foreachIncludeOfStoreUnitContainingFile(CanonicalFilePathRef filePath,
                    function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver) {
   SmallVector<std::string, 32> allUnitNames;
@@ -241,19 +261,7 @@ bool FileIndexImpl::foreachIncludeOfStoreUnitContainingFile(CanonicalFilePathRef
   }
 
   for (auto &unitName: allUnitNames) {
-    std::string error;
-    IndexUnitReader storeUnit(*IdxStore, unitName, error);
-    if (!storeUnit.isValid())
-      continue;
-    StringRef workDir = storeUnit.getWorkingDirectory();
-    bool cont = storeUnit.foreachInclude([&](IndexUnitInclude inc) -> bool {
-      StringRef sourcePath = inc.getSourcePath();
-      StringRef targetPath = inc.getTargetPath();
-      unsigned line = inc.getSourceLine();
-      CanonicalFilePath fullSourcePath = getCanonicalPath(sourcePath, workDir);
-      CanonicalFilePath fullTargetPath = getCanonicalPath(targetPath, workDir);
-      return receiver(fullSourcePath, fullTargetPath, line);
-    });
+    bool cont = foreachIncludeOfUnit(unitName, receiver);
     if (!cont)
       return false;
   }
@@ -314,3 +322,7 @@ bool FilePathIndex::foreachFileIncludingFile(CanonicalFilePathRef TargetPath, fu
 bool FilePathIndex::foreachFileIncludedByFile(CanonicalFilePathRef SourcePath, function_ref<bool (CanonicalFilePathRef, unsigned int)> Receiver) {
   return IMPL->foreachFileIncludedByFile(SourcePath, Receiver);
 }
+
+bool FilePathIndex::foreachIncludeOfUnit(StringRef unitName, function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver) {
+  return IMPL->foreachIncludeOfUnit(unitName, receiver);
+}

lib/Index/IndexDatastore.cpp

@@ -573,7 +573,7 @@ void StoreUnitRepo::registerUnit(StringRef unitName, std::shared_ptr<UnitProcess
           if (moduleName.empty()) {
             // Workaround for swift compiler not associating the module name with records of swift files.
             // FIXME: Fix this on swift compiler and remove this.
-            if (CanonPath.getPath().endswith(".swift")) {
+            if (StringRef(CanonPath.getPath()).endswith(".swift")) {
               moduleName = Reader.getModuleName();
             }
           }

lib/Index/IndexSystem.cpp

@@ -195,6 +195,9 @@ class IndexSystemImpl {
   bool foreachFileIncludedByFile(StringRef SourcePath,
                                  function_ref<bool(CanonicalFilePathRef TargetPath, unsigned Line)> Receiver);
 
+  bool foreachIncludeOfUnit(StringRef unitName,
+                            function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver);
+
   bool foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths,
       function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 };
@@ -565,6 +568,11 @@ bool IndexSystemImpl::foreachFileIncludedByFile(StringRef SourcePath,
   return PathIndex->foreachFileIncludedByFile(canonSourcePath, Receiver);
 }
 
+bool IndexSystemImpl::foreachIncludeOfUnit(StringRef unitName,
+                                           function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver) {
+  return PathIndex->foreachIncludeOfUnit(unitName, receiver);
+}
+
 bool IndexSystemImpl::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths, function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
   return SymIndex->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));
 }
@@ -764,6 +772,11 @@ bool IndexSystem::foreachFileIncludedByFile(StringRef SourcePath,
   return IMPL->foreachFileIncludedByFile(SourcePath, Receiver);
 }
 
+bool IndexSystem::foreachIncludeOfUnit(StringRef unitName,
+                                       function_ref<bool(CanonicalFilePathRef sourcePath, CanonicalFilePathRef targetPath, unsigned line)> receiver) {
+  return IMPL->foreachIncludeOfUnit(unitName, receiver);
+}
+
 bool IndexSystem::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths,
     function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
   return IMPL->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));