Symbol occurrences in file

.gitignore

@@ -1,4 +1,5 @@
 .DS_Store
+.vscode
 default.profraw
 Package.resolved
 /.build

Sources/IndexStoreDB/IndexStoreDB.swift

@@ -312,6 +312,41 @@ public final class IndexStoreDB {
     return result
   }
 
+  public func symbolProvider(for sourceFilePath: String) -> SymbolProviderKind? {
+    var result: SymbolProviderKind? = nil
+    indexstoredb_index_units_containing_file(impl, sourceFilePath) { unit in
+      let providerKind: SymbolProviderKind? = switch indexstoredb_unit_info_symbol_provider_kind(unit) {
+      case INDEXSTOREDB_SYMBOL_PROVIDER_KIND_SWIFT:
+        .swift
+      case INDEXSTOREDB_SYMBOL_PROVIDER_KIND_CLANG:
+        .clang
+      case INDEXSTOREDB_SYMBOL_PROVIDER_KIND_UNKNOWN:
+        nil
+      default:
+        preconditionFailure("Unknown enum case in indexstoredb_symbol_provider_kind_t")
+      }
+
+      let mainFilePath = String(cString: indexstoredb_unit_info_main_file_path(unit))
+      if providerKind == .swift && mainFilePath != sourceFilePath {
+        // We have a unit that is "included" from Swift. This happens for header
+        // files that Swift files depend on. But Swift doesn't have includes and
+        // we shouldn't infer the header file's language to Swift based on this unit.
+        // Ignore it.
+        return true
+      }
+
+      if result == nil {
+        result = providerKind
+      } else if result != providerKind {
+        // Found two conflicting provider kinds. Return nil as we don't know the provider in this case.
+        result = nil
+        return false
+      }
+      return true
+    }
+    return result
+  }
+
   @discardableResult
   public func foreachFileIncludedByFile(path: String, body: (String) -> Bool) -> Bool {
     return withoutActuallyEscaping(body) { body in
@@ -321,7 +356,7 @@ public final class IndexStoreDB {
       }
     }
   }
-    
+
   public func filesIncludedByFile(path: String) -> [String] {
     var result: [String] = []
     foreachFileIncludedByFile(path: path) { targetPath in
@@ -330,7 +365,7 @@ public final class IndexStoreDB {
     }
     return result
   }
-    
+
   @discardableResult
   public func foreachFileIncludingFile(path: String, body: (String) -> Bool) -> Bool {
     return withoutActuallyEscaping(body) { body in
@@ -340,7 +375,7 @@ public final class IndexStoreDB {
       }
     }
   }
-      
+
   public func filesIncludingFile(path: String) -> [String] {
     var result: [String] = []
     foreachFileIncludingFile(path: path) { targetPath in
@@ -421,6 +456,15 @@ public final class IndexStoreDB {
     return result
   }
 
+  public func symbolOccurrences(inFilePath path: String) -> [SymbolOccurrence] {
+    var result: [SymbolOccurrence] = []
+    forEachSymbolOccurrence(inFilePath: path) { occur in
+      result.append(occur)
+      return true
+    }
+    return result
+  }
+
   @discardableResult
   func forEachSymbol(inFilePath filePath: String, body: (Symbol) -> Bool) -> Bool {
     return withoutActuallyEscaping(body) { body in
@@ -430,6 +474,15 @@ public final class IndexStoreDB {
     }
   }
 
+  @discardableResult
+  func forEachSymbolOccurrence(inFilePath filePath: String, body: (SymbolOccurrence) -> Bool) -> Bool {
+    return withoutActuallyEscaping(body) { body in
+      return indexstoredb_index_symbol_occurrences_in_file_path(impl, filePath) { occur in
+        return body(SymbolOccurrence(occur))
+      }
+    }
+  }
+
   /// 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] = []

include/CIndexStoreDB/CIndexStoreDB.h

@@ -328,6 +328,15 @@ indexstoredb_index_symbols_contained_in_file_path(_Nonnull indexstoredb_index_t
                                                   const char *_Nonnull path,
                                                   _Nonnull indexstoredb_symbol_receiver_t);
 
+/// Iterates over all the symbol occurrences contained in \p path
+///
+/// The occurrence passed to the receiver is only valid for the duration of the
+/// receiver call.
+INDEXSTOREDB_PUBLIC bool
+indexstoredb_index_symbol_occurrences_in_file_path(_Nonnull indexstoredb_index_t index,
+                                                   const char *_Nonnull path,
+                                                   _Nonnull indexstoredb_symbol_occurrence_receiver_t);
+
 /// Returns the USR of the given symbol.
 ///
 /// The string has the same lifetime as the \c indexstoredb_symbol_t.

include/IndexStoreDB/Index/IndexSystem.h

@@ -104,6 +104,9 @@ class INDEXSTOREDB_EXPORT IndexSystem {
   bool foreachSymbolInFilePath(StringRef FilePath,
                                function_ref<bool(SymbolRef Symbol)> Receiver);
 
+  bool foreachSymbolOccurrenceInFilePath(StringRef FilePath,
+                                         function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+
   bool foreachSymbolOccurrenceByUSR(StringRef USR, SymbolRoleSet RoleSet,
                         function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 
@@ -178,7 +181,7 @@ class INDEXSTOREDB_EXPORT IndexSystem {
   bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
 
   /// Returns the latest modification date of a unit that contains the given source file.
-  /// 
+  ///
   /// If no unit containing the given source file exists, returns `None`.
   llvm::Optional<llvm::sys::TimePoint<>> timestampOfLatestUnitForFile(StringRef filePath);
 private:

include/IndexStoreDB/Index/SymbolDataProvider.h

@@ -48,6 +48,8 @@ class SymbolDataProvider {
                                                        SymbolRoleSet Roles,
                                                        SymbolRoleSet RelatedRoles)> Receiver) = 0;
 
+  virtual bool foreachSymbolOccurrence(function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) = 0;
+
   virtual bool foreachSymbolOccurrenceByUSR(ArrayRef<db::IDCode> USRs,
                                             SymbolRoleSet RoleSet,
                         function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) = 0;

include/IndexStoreDB/Index/SymbolIndex.h

@@ -68,6 +68,9 @@ class SymbolIndex {
   bool foreachSymbolInFilePath(CanonicalFilePathRef filePath,
                                function_ref<bool(SymbolRef Occur)> Receiver);
 
+  bool foreachSymbolOccurrenceInFilePath(CanonicalFilePathRef filePath,
+                                         function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+
   bool foreachCanonicalSymbolOccurrenceContainingPattern(StringRef Pattern,
                                                 bool AnchorStart,
                                                 bool AnchorEnd,
@@ -104,7 +107,7 @@ class SymbolIndex {
   bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
 
   /// Returns the latest modification date of a unit that contains the given source file.
-  /// 
+  ///
   /// If no unit containing the given source file exists, returns `None`.
   llvm::Optional<llvm::sys::TimePoint<>> timestampOfLatestUnitForFile(CanonicalFilePathRef filePath);
 private:

lib/CIndexStoreDB/CIndexStoreDB.cpp

@@ -314,6 +314,16 @@ indexstoredb_index_symbols_contained_in_file_path(_Nonnull indexstoredb_index_t
   });
 }
 
+bool
+indexstoredb_index_symbol_occurrences_in_file_path(_Nonnull indexstoredb_index_t index,
+                                                    const char *_Nonnull path,
+                                                    _Nonnull indexstoredb_symbol_occurrence_receiver_t receiver) {
+  auto obj = (Object<std::shared_ptr<IndexSystem>> *)index;
+  return obj->value->foreachSymbolOccurrenceInFilePath(path, [&](SymbolOccurrenceRef Occur) -> bool {
+    return receiver((indexstoredb_symbol_occurrence_t)Occur.get());
+  });
+}
+
 const char *
 indexstoredb_symbol_usr(indexstoredb_symbol_t symbol) {
   auto value = (Symbol *)symbol;

lib/Index/IndexSystem.cpp

@@ -160,11 +160,14 @@ class IndexSystemImpl {
 
   void addDelegate(std::shared_ptr<IndexSystemDelegate> Delegate);
 
-    bool foreachSymbolInFilePath(StringRef filePath,
-                                 function_ref<bool(const SymbolRef &symbol)> receiver);
-
-    bool foreachSymbolOccurrenceByUSR(StringRef USR, SymbolRoleSet RoleSet,
-                        function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+  bool foreachSymbolInFilePath(StringRef filePath,
+                               function_ref<bool(const SymbolRef &symbol)> receiver);
+
+  bool foreachSymbolOccurrenceInFilePath(StringRef filePath,
+                                         function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+
+  bool foreachSymbolOccurrenceByUSR(StringRef USR, SymbolRoleSet RoleSet,
+                                    function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
 
   bool foreachCanonicalSymbolOccurrenceContainingPattern(StringRef Pattern,
                                                 bool AnchorStart,
@@ -236,7 +239,7 @@ class IndexSystemImpl {
   bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
 
   /// Returns the latest modification date of a unit that contains the given source file.
-  /// 
+  ///
   /// If no unit containing the given source file exists, returns `None`.
   llvm::Optional<llvm::sys::TimePoint<>> timestampOfLatestUnitForFile(StringRef filePath);
 };
@@ -579,6 +582,12 @@ bool IndexSystemImpl::foreachSymbolInFilePath(StringRef filePath,
     return SymIndex->foreachSymbolInFilePath(canonPath, std::move(receiver));
 }
 
+bool IndexSystemImpl::foreachSymbolOccurrenceInFilePath(StringRef filePath,
+                                                        function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  auto canonPath = PathIndex->getCanonicalPath(filePath);
+  return SymIndex->foreachSymbolOccurrenceInFilePath(canonPath, std::move(Receiver));
+}
+
 bool IndexSystemImpl::foreachFileOfUnit(StringRef unitName,
                                         bool followDependencies,
                                         function_ref<bool(CanonicalFilePathRef filePath)> receiver) {
@@ -778,6 +787,11 @@ bool IndexSystem::foreachSymbolInFilePath(StringRef FilePath,
     return IMPL->foreachSymbolInFilePath(FilePath, std::move(Receiver));
 }
 
+bool IndexSystem::foreachSymbolOccurrenceInFilePath(StringRef FilePath,
+                                                    function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  return IMPL->foreachSymbolOccurrenceInFilePath(FilePath, std::move(Receiver));
+}
+
 bool IndexSystem::isKnownFile(StringRef filePath) {
   return IMPL->isKnownFile(filePath);
 }

lib/Index/StoreSymbolRecord.cpp

@@ -289,6 +289,20 @@ bool StoreSymbolRecord::foreachCoreSymbolData(function_ref<bool(StringRef USR,
   return !Err && Finished;
 }
 
+bool StoreSymbolRecord::foreachSymbolOccurrence(function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  bool Finished;
+  bool Err = doForData([&](IndexRecordReader &Reader) {
+      // Return all occurrences.
+    auto Pred = [](IndexRecordOccurrence) -> bool { return true; };
+    PredOccurrenceConverter Converter(*this, Pred, Receiver);
+    Finished = Reader.foreachOccurrence(/*symbolsFilter=*/None,
+                                        /*relatedSymbolsFilter=*/None,
+                                        Converter);
+  });
+
+  return !Err && Finished;
+}
+
 static void searchDeclsByUSR(IndexRecordReader &Reader,
                              ArrayRef<db::IDCode> USRs,
            SmallVectorImpl<IndexRecordSymbol> &FoundDecls) {

lib/Index/StoreSymbolRecord.h

@@ -75,6 +75,8 @@ class StoreSymbolRecord : public SymbolDataProvider {
                                                        SymbolRoleSet Roles,
                                                        SymbolRoleSet RelatedRoles)> Receiver) override;
 
+  virtual bool foreachSymbolOccurrence(function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) override;
+
   virtual bool foreachSymbolOccurrenceByUSR(ArrayRef<db::IDCode> USRs,
                                             SymbolRoleSet RoleSet,
                function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) override;

lib/Index/SymbolIndex.cpp

@@ -76,6 +76,9 @@ class SymbolIndexImpl {
   bool foreachSymbolInFilePath(CanonicalFilePathRef filePath,
                                function_ref<bool(SymbolRef Symbol)> Receiver);
 
+  bool foreachSymbolOccurrenceInFilePath(CanonicalFilePathRef filePath,
+                                         function_ref<bool(SymbolOccurrenceRef Occur)> Receiver);
+
   bool foreachSymbolName(function_ref<bool(StringRef name)> receiver);
 
   bool foreachCanonicalSymbolOccurrenceByUSR(StringRef USR,
@@ -100,7 +103,7 @@ class SymbolIndexImpl {
   bool foreachUnitTestSymbol(function_ref<bool(SymbolOccurrenceRef Occur)> receiver);
 
   /// Returns the latest modification date of a unit that contains the given source file.
-  /// 
+  ///
   /// If no unit containing the given source file exists, returns `None`.
   llvm::Optional<sys::TimePoint<>> timestampOfLatestUnitForFile(CanonicalFilePathRef filePath);
 
@@ -481,7 +484,7 @@ bool SymbolIndexImpl::foreachSymbolInFilePath(CanonicalFilePathRef filePath,
                                                                   SymbolInfo info,
                                                                   SymbolRoleSet roles,
                                                                   SymbolRoleSet relatedRoles) -> bool {
-                      
+
                       if (roles.containsAny(SymbolRoleSet(SymbolRole::Definition) | SymbolRole::Declaration)) {
                         return Receiver(std::make_shared<Symbol>(info, name, usr));
                       } else {
@@ -500,6 +503,33 @@ bool SymbolIndexImpl::foreachSymbolInFilePath(CanonicalFilePathRef filePath,
     return didFinish;
 }
 
+bool SymbolIndexImpl::foreachSymbolOccurrenceInFilePath(CanonicalFilePathRef filePath,
+                                                        function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+    bool didFinish = true;
+    ReadTransaction reader(DBase);
+
+    IDCode filePathCode = reader.getFilePathCode(filePath);
+    reader.foreachUnitContainingFile(filePathCode, [&](ArrayRef<IDCode> idCodes) -> bool {
+        for (IDCode idCode : idCodes) {
+            UnitInfo unitInfo = reader.getUnitInfo(idCode);
+
+            for (UnitInfo::Provider provider : unitInfo.ProviderDepends) {
+                IDCode providerCode = provider.ProviderCode;
+                if (provider.FileCode == filePathCode) {
+                    auto record = createVisibleProviderForCode(providerCode, reader);
+                    didFinish = record->foreachSymbolOccurrence(Receiver);
+
+                    return false;
+                }
+            }
+        }
+
+        return true;
+    });
+
+    return didFinish;
+}
+
 bool SymbolIndexImpl::foreachCanonicalSymbolOccurrenceByKind(SymbolKind symKind, bool workspaceOnly,
                                                              function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
   return foreachCanonicalSymbolOccurrenceImpl(workspaceOnly,
@@ -706,6 +736,11 @@ bool SymbolIndex::foreachSymbolInFilePath(CanonicalFilePathRef filePath,
   return IMPL->foreachSymbolInFilePath(filePath, std::move(Receiver));
 }
 
+bool SymbolIndex::foreachSymbolOccurrenceInFilePath(CanonicalFilePathRef filePath,
+                                                    function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
+  return IMPL->foreachSymbolOccurrenceInFilePath(filePath, std::move(Receiver));
+}
+
 bool SymbolIndex::foreachUnitTestSymbolReferencedByOutputPaths(ArrayRef<CanonicalFilePathRef> FilePaths,
     function_ref<bool(SymbolOccurrenceRef Occur)> Receiver) {
   return IMPL->foreachUnitTestSymbolReferencedByOutputPaths(FilePaths, std::move(Receiver));