Merge pull request #77373 from nkcsgexi/env-var-objc-message

ModuleObjCTrace: allow configuring the directory path to where the objc message trace file will be emitted to via an environment variable

Author: nkcsgexi

lib/FrontendTool/LoadedModuleTrace.cpp

@@ -910,30 +910,45 @@ class ObjcMethodReferenceCollector: public SourceEntityWalker {
   }
 };
 
+static std::optional<int> createObjCMessageTraceFile(const InputFile &input,
+                                                     ModuleDecl *MD) {
+  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()) {
+    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))
+      return {};
+  }
+  SmallString<32> fileName(MD->getNameStr());
+  fileName.append("-%%%%-%%%%-%%%%.json");
+  llvm::sys::path::append(tracePath, fileName);
+  int tmpFD;
+  if (llvm::sys::fs::createUniqueFile(tracePath.str(), tmpFD, tracePath)) {
+    return {};
+  }
+  return tmpFD;
+}
+
 bool swift::emitObjCMessageSendTraceIfNeeded(ModuleDecl *mainModule,
                                              const FrontendOptions &opts) {
   ASTContext &ctxt = mainModule->getASTContext();
   assert(!ctxt.hadError() &&
          "We should've already exited earlier if there was an error.");
 
   opts.InputsAndOutputs.forEachInput([&](const InputFile &input) {
-    auto loadedModuleTracePath = input.getLoadedModuleTracePath();
-    if (loadedModuleTracePath.empty())
-      return false;
-    llvm::SmallString<128> tracePath {loadedModuleTracePath};
-    llvm::sys::path::remove_filename(tracePath);
-    llvm::sys::path::append(tracePath, ".SWIFT_FINE_DEPENDENCY_TRACE");
-    if (!llvm::sys::fs::exists(tracePath)) {
-      if (llvm::sys::fs::create_directory(tracePath))
-        return false;
-    }
-    llvm::sys::path::append(tracePath, "%%%%-%%%%-%%%%.json");
-    int tmpFD;
-    if (llvm::sys::fs::createUniqueFile(tracePath.str(), tmpFD, tracePath)) {
+    auto tmpFD = createObjCMessageTraceFile(input, mainModule);
+    if (!tmpFD)
       return false;
-    }
     // Write the contents of the buffer.
-    llvm::raw_fd_ostream out(tmpFD, /*shouldClose=*/true);
+    llvm::raw_fd_ostream out(*tmpFD, /*shouldClose=*/true);
     ObjcMethodReferenceCollector collector(mainModule);
     for (auto *FU : mainModule->getFiles()) {
       if (auto *SF = dyn_cast<SourceFile>(FU)) {

test/IDE/objc_send_collector_1.swift

@@ -1,8 +1,12 @@
 // RUN: %empty-directory(%t)
+// RUN: %empty-directory(%t/CUSTOM_DIR)
 
 // 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: cat %t/CUSTOM_DIR/* | %FileCheck %s
+
 // REQUIRES: objc_interop
 
 import Foo