ModuleObjCTrace: walk only Swift source files to collect used ObjC method invocations

Author: nkcsgexi

lib/FrontendTool/LoadedModuleTrace.cpp

@@ -911,17 +911,34 @@ class ObjcMethodReferenceCollector: public SourceEntityWalker {
 };
 
 static std::optional<int> createObjCMessageTraceFile(const InputFile &input,
-                                                     ModuleDecl *MD) {
+                                                     ModuleDecl *MD,
+                                        std::vector<SourceFile*> &filesToWalk) {
+  if (input.getLoadedModuleTracePath().empty()) {
+    // we basically rely on the passing down of module trace file path
+    // as an indicator that this job needs to emit an ObjC message trace file.
+    // FIXME: add a separate swift-frontend flag for ObjC message trace path
+    // specifically.
+    return {};
+  }
+  for (auto *FU : MD->getFiles()) {
+    if (auto *SF = dyn_cast<SourceFile>(FU)) {
+      if (SF->getFilename().ends_with(".swift")) {
+        filesToWalk.push_back(SF);
+      }
+    }
+  }
+  // No source files to walk, abort.
+  if (filesToWalk.empty()) {
+    return {};
+  }
   llvm::SmallString<128> tracePath;
   if (const char *P = ::getenv("SWIFT_COMPILER_OBJC_MESSAGE_TRACE_DIRECTORY")) {
     StringRef DirPath = P;
     llvm::sys::path::append(tracePath, DirPath);
-  } else if (!input.getLoadedModuleTracePath().empty()) {
+  } else {
     llvm::sys::path::append(tracePath, input.getLoadedModuleTracePath());
     llvm::sys::path::remove_filename(tracePath);
     llvm::sys::path::append(tracePath, ".SWIFT_FINE_DEPENDENCY_TRACE");
-  } else {
-    return {};
   }
   if (!llvm::sys::fs::exists(tracePath)) {
     if (llvm::sys::fs::create_directory(tracePath))
@@ -944,17 +961,16 @@ bool swift::emitObjCMessageSendTraceIfNeeded(ModuleDecl *mainModule,
          "We should've already exited earlier if there was an error.");
 
   opts.InputsAndOutputs.forEachInput([&](const InputFile &input) {
-    auto tmpFD = createObjCMessageTraceFile(input, mainModule);
+    std::vector<SourceFile*> filesToWalk;
+    auto tmpFD = createObjCMessageTraceFile(input, mainModule, filesToWalk);
     if (!tmpFD)
       return false;
     // Write the contents of the buffer.
     llvm::raw_fd_ostream out(*tmpFD, /*shouldClose=*/true);
     ObjcMethodReferenceCollector collector(mainModule);
-    for (auto *FU : mainModule->getFiles()) {
-      if (auto *SF = dyn_cast<SourceFile>(FU)) {
-        collector.setFileBeforeVisiting(SF);
-        collector.walk(*SF);
-      }
+    for (auto *SF : filesToWalk) {
+      collector.setFileBeforeVisiting(SF);
+      collector.walk(*SF);
     }
     collector.serializeAsJson(out);
     return true;

test/IDE/objc_send_collector_1.swift

@@ -4,7 +4,7 @@
 // RUN: %target-swift-frontend  -I %t/lib/swift -typecheck %s %S/Inputs/objc_send_collector_2.swift -module-name main -swift-version 5 -F %S/Inputs/mock-sdk -emit-loaded-module-trace-path %t/.MODULE_TRACE
 // RUN: cat %t/.SWIFT_FINE_DEPENDENCY_TRACE/* | %FileCheck %s
 
-// RUN: SWIFT_COMPILER_OBJC_MESSAGE_TRACE_DIRECTORY=%t/CUSTOM_DIR %target-swift-frontend  -I %t/lib/swift -typecheck %s %S/Inputs/objc_send_collector_2.swift -module-name main -swift-version 5 -F %S/Inputs/mock-sdk
+// RUN: SWIFT_COMPILER_OBJC_MESSAGE_TRACE_DIRECTORY=%t/CUSTOM_DIR %target-swift-frontend  -I %t/lib/swift -typecheck %s %S/Inputs/objc_send_collector_2.swift -module-name main -swift-version 5 -F %S/Inputs/mock-sdk -emit-loaded-module-trace-path %t/.MODULE_TRACE
 // RUN: cat %t/CUSTOM_DIR/* | %FileCheck %s
 
 // REQUIRES: objc_interop