Make sure Swift symbol occurrences are associated with the module name they belong to

Author: akyrtzi

Sources/ISDBTestSupport/TestLocation.swift

@@ -16,6 +16,9 @@ import IndexStoreDB
 /// A source location (file:line:column) in a test project, for use with the TestLocationScanner.
 public struct TestLocation: Hashable {
 
+  /// Marker string to mark a `SymbolLocation` as containing an unknown module name from a test location.
+  public static let unknownModuleName: String = "<UNKNOWN>"
+
   /// The path/url of the source file.
   public var url: URL
 
@@ -45,9 +48,10 @@ extension TestLocation: Comparable {
 extension SymbolLocation {
 
   /// Constructs a SymbolLocation from a TestLocation, using a non-system path by default.
-  public init(_ loc: TestLocation, isSystem: Bool = false) {
+  public init(_ loc: TestLocation, moduleName: String = TestLocation.unknownModuleName, isSystem: Bool = false) {
     self.init(
       path: loc.url.path,
+      moduleName: moduleName,
       isSystem: isSystem,
       line: loc.line,
       utf8Column: loc.utf8Column)
@@ -57,8 +61,8 @@ extension SymbolLocation {
 extension Symbol {
 
   /// Returns a SymbolOccurrence with the given location and roles.
-  public func at(_ location: TestLocation, roles: SymbolRole) -> SymbolOccurrence {
-    return self.at(SymbolLocation(location), roles: roles)
+  public func at(_ location: TestLocation, moduleName: String = TestLocation.unknownModuleName, roles: SymbolRole) -> SymbolOccurrence {
+    return self.at(SymbolLocation(location, moduleName: moduleName), roles: roles)
   }
 }
 

Sources/ISDBTestSupport/XCTestAssertions.swift

@@ -13,6 +13,32 @@
 import IndexStoreDB
 import XCTest
 
+extension SymbolLocation {
+  /// Special equality check for testing.
+  func testEqual(with other: SymbolLocation) -> Bool {
+    // We do not have a way to automatically infer the module name from test locations,
+    // if we used `SymbolLocation`'s standard equality check we'd have to manually pass
+    // the module name string in every symbol occurrence test.
+    // Instead we ignore the module name if this `SymbolLocation` was initialized from
+    // a test location without an explicit module name.
+    return path == other.path &&
+      (moduleName == TestLocation.unknownModuleName || other.moduleName == TestLocation.unknownModuleName || moduleName == other.moduleName) &&
+      isSystem == other.isSystem &&
+      line == other.line &&
+      utf8Column == other.utf8Column
+  }
+}
+
+extension SymbolOccurrence {
+  /// Special equality check for testing.
+  func testEqual(with other: SymbolOccurrence) -> Bool {
+    return symbol == other.symbol &&
+      location.testEqual(with: other.location) &&
+      roles == other.roles &&
+      relations == other.relations
+  }
+}
+
 /// An XCTest assertion that the given symbol occurrences match the expected values, ignoring the
 /// order of results, and optionally ignoring relations. By default, the given occurrences also are
 /// allowed to have additional SymbolRoles.
@@ -59,7 +85,7 @@ public func checkOccurrences(
       actual.roles.formIntersection(expected.roles)
     }
 
-    if actual == expected { return .orderedSame }
+    if actual.testEqual(with: expected) { return .orderedSame }
     if actual < expected { return .orderedAscending }
     return .orderedDescending
   }

Sources/IndexStoreDB/SymbolLocation.swift

@@ -14,12 +14,14 @@ import CIndexStoreDB
 
 public struct SymbolLocation: Equatable {
   public var path: String
+  public var moduleName: String
   public var isSystem: Bool
   public var line: Int
   public var utf8Column: Int
 
-  public init(path: String, isSystem: Bool = false, line: Int, utf8Column: Int) {
+  public init(path: String, moduleName: String, isSystem: Bool = false, line: Int, utf8Column: Int) {
     self.path = path
+    self.moduleName = moduleName
     self.isSystem = isSystem
     self.line = line
     self.utf8Column = utf8Column
@@ -28,14 +30,14 @@ public struct SymbolLocation: Equatable {
 
 extension SymbolLocation: Comparable {
   public static func <(a: SymbolLocation, b: SymbolLocation) -> Bool {
-    return (a.path, a.line, a.utf8Column, a.isSystem ? 1 : 0)
-      < (b.path, b.line, b.utf8Column, b.isSystem ? 1 : 0)
+    return (a.path, a.moduleName, a.line, a.utf8Column, a.isSystem ? 1 : 0)
+      < (b.path, a.moduleName, b.line, b.utf8Column, b.isSystem ? 1 : 0)
   }
 }
 
 extension SymbolLocation: CustomStringConvertible {
   public var description: String {
-    "\(path):\(line):\(utf8Column)\(isSystem ? " [system]" : "")"
+    "\(!moduleName.isEmpty ? "\(moduleName)::" : "")\(path):\(line):\(utf8Column)\(isSystem ? " [system]" : "")"
   }
 }
 
@@ -44,6 +46,7 @@ extension SymbolLocation: CustomStringConvertible {
 extension SymbolLocation {
   public init(_ loc: indexstoredb_symbol_location_t) {
     path = String(cString: indexstoredb_symbol_location_path(loc))
+    moduleName = String(cString: indexstoredb_symbol_location_module_name(loc))
     isSystem = indexstoredb_symbol_location_is_system(loc)
     line = Int(indexstoredb_symbol_location_line(loc))
     utf8Column = Int(indexstoredb_symbol_location_column_utf8(loc))

Tests/IndexStoreDBTests/IndexTests.swift

@@ -32,7 +32,7 @@ final class IndexTests: XCTestCase {
 
     let ccanon = SymbolOccurrence(
       symbol: csym,
-      location: SymbolLocation(ws.testLoc("c")),
+      location: SymbolLocation(ws.testLoc("c"), moduleName: "main"),
       roles: [.definition, .canonical],
       relations: [])
 
@@ -135,9 +135,9 @@ final class IndexTests: XCTestCase {
 
     let aaa = Symbol(usr: "s:1A3aaayyF", name: "aaa()", kind: .function)
     checkOccurrences(ws.index.occurrences(ofUSR: aaa.usr, roles: .all), expected: [
-      aaa.at(ws.testLoc("aaa:def"), roles: .definition),
-      aaa.at(ws.testLoc("aaa:call"), roles: .call),
-      aaa.at(ws.testLoc("aaa:call:c"), roles: .call),
+      aaa.at(ws.testLoc("aaa:def"), moduleName: "A", roles: .definition),
+      aaa.at(ws.testLoc("aaa:call"), moduleName: "B", roles: .call),
+      aaa.at(ws.testLoc("aaa:call:c"), moduleName: "C", roles: .call),
     ])
   }
 

include/CIndexStoreDB/CIndexStoreDB.h

@@ -231,6 +231,13 @@ INDEXSTOREDB_PUBLIC
 const char * _Nonnull
 indexstoredb_symbol_location_path(_Nonnull indexstoredb_symbol_location_t);
 
+/// Returns the module name of the given symbol location.
+///
+/// The string has the same lifetime as the \c indexstoredb_symbol_location_t.
+INDEXSTOREDB_PUBLIC
+const char * _Nonnull
+indexstoredb_symbol_location_module_name(_Nonnull indexstoredb_symbol_location_t);
+
 /// Returns whether the given symbol location is a system location.
 INDEXSTOREDB_PUBLIC bool
 indexstoredb_symbol_location_is_system(_Nonnull indexstoredb_symbol_location_t);

lib/CIndexStoreDB/CIndexStoreDB.cpp

@@ -286,6 +286,12 @@ indexstoredb_symbol_location_path(indexstoredb_symbol_location_t loc) {
   return obj->getPath().getPathString().c_str();
 }
 
+const char *
+indexstoredb_symbol_location_module_name(indexstoredb_symbol_location_t loc) {
+  auto obj = (SymbolLocation *)loc;
+  return obj->getPath().getModuleName().c_str();
+}
+
 bool
 indexstoredb_symbol_location_is_system(indexstoredb_symbol_location_t loc) {
   auto obj = (SymbolLocation *)loc;

lib/Index/IndexDatastore.cpp

@@ -367,8 +367,16 @@ void StoreUnitRepo::registerUnit(StringRef unitName) {
           if (!Dep.isSystem())
             UserFileDepends.push_back(CanonPath);
           StringRef recordName = Dep.getName();
+          StringRef moduleName = Dep.getModuleName();
+          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")) {
+              moduleName = Reader.getModuleName();
+            }
+          }
           bool isNewProvider;
-          IDCode providerCode = unitImport.addProviderDependency(recordName, CanonPath, Dep.getModuleName(), Dep.isSystem(), &isNewProvider);
+          IDCode providerCode = unitImport.addProviderDependency(recordName, CanonPath, moduleName, Dep.isSystem(), &isNewProvider);
           if (!isNewProvider)
             break;