From 94d133307223810ad5191bd16225958f0f75ec32 Mon Sep 17 00:00:00 2001
From: Xi Ge <xi_ge@apple.com>
Date: Mon, 4 Nov 2024 18:21:37 -0800
Subject: [PATCH] ModuleObjCTrace: walk only Swift source files to collect used
 ObjC method invocations

---
 lib/FrontendTool/LoadedModuleTrace.cpp | 36 +++++++++++++++++++-------
 test/IDE/objc_send_collector_1.swift   |  2 +-
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/lib/FrontendTool/LoadedModuleTrace.cpp b/lib/FrontendTool/LoadedModuleTrace.cpp
index e3820a42ac227..3930fa8aa5d9d 100644
--- a/lib/FrontendTool/LoadedModuleTrace.cpp
+++ b/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;
diff --git a/test/IDE/objc_send_collector_1.swift b/test/IDE/objc_send_collector_1.swift
index d225e1de9c92d..e354817ba0276 100644
--- a/test/IDE/objc_send_collector_1.swift
+++ b/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