[llvm] [llvm-driver] Fix tool re-entrance on Windows. (PR #76306)

[Alexandre Ganea via llvm-commits]

https://github.com/aganea created https://github.com/llvm/llvm-project/pull/76306

Previously, some tools such as `clang` or `lld` which require strict order for certain command-line options, such as `clang -cc1` or `lld -flavor` would not long work on Windows, when these tools were linked as part of `llvm-driver`. This was caused by `InitLLVM` which was part of the `main()` function of these tools, which in turn calls `windows::GetCommandLineArguments`. This function completly replaces argc/argv by new UTF-8 contents, so any ajustements to argc/argv made by `llvm-driver` prior to calling these tools would be reset.

We now call `InitLLVM` as part of the `llvm-driver`. Any further usages to `InitLLVM` on the stack, after the first call in the process would have no effect. In the same way, the last `InitLLVM` on the stack will clear the `ManagedStatics` as usual.

>From 8187e95ea4e04793fbfc85045aa21f9633bbc03d Mon Sep 17 00:00:00 2001
From: Alexandre Ganea <alex_toresh at yahoo.fr>
Date: Sat, 23 Dec 2023 19:17:06 -0500
Subject: [PATCH] [llvm-driver] Fix tool re-entrance on Windows.

Previously, some tools such as `clang` or `lld` which require strict order for certain command-line options, such as `clang -cc1` or `lld -flavor` would not long work on Windows, when these tools were linked as part of `llvm-driver`. This was caused by `InitLLVM` which was part of the `main()` function of these tools, which in turn calls `windows::GetCommandLineArguments`. This function completly replaces argc/argv by new UTF-8 contents, so any ajustements to argc/argv made by `llvm-driver` prior to calling these tools would be reset.

We now call `InitLLVM` as part of the `llvm-driver`. Any further usages to `InitLLVM` on the stack, after the first call in the process would have no effect. In the same way, the last `InitLLVM` on the stack will clear the `ManagedStatics` as usual.
---
 llvm/cmake/modules/llvm-driver-template.cpp.in | 2 ++
 llvm/lib/Support/InitLLVM.cpp                  | 6 ++++++
 llvm/tools/llvm-driver/llvm-driver.cpp         | 6 +++++-
 3 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/llvm/cmake/modules/llvm-driver-template.cpp.in b/llvm/cmake/modules/llvm-driver-template.cpp.in
index 16c4fb34714638..71aca6cd140cb5 100644
--- a/llvm/cmake/modules/llvm-driver-template.cpp.in
+++ b/llvm/cmake/modules/llvm-driver-template.cpp.in
@@ -8,9 +8,11 @@
 
 #include "llvm/Support/LLVMDriver.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/InitLLVM.h"
 
 int @TOOL_NAME at _main(int argc, char **, const llvm::ToolContext &);
 
 int main(int argc, char **argv) {
+  llvm::InitLLVM X(argc, argv);
   return @TOOL_NAME at _main(argc, argv, {argv[0], nullptr, false});
 }
diff --git a/llvm/lib/Support/InitLLVM.cpp b/llvm/lib/Support/InitLLVM.cpp
index 7f475f42f3cb81..2a2e6c254c795a 100644
--- a/llvm/lib/Support/InitLLVM.cpp
+++ b/llvm/lib/Support/InitLLVM.cpp
@@ -36,8 +36,12 @@ void CleanupStdHandles(void *Cookie) {
 using namespace llvm;
 using namespace llvm::sys;
 
+static std::atomic<unsigned> UsageCount{0};
+
 InitLLVM::InitLLVM(int &Argc, const char **&Argv,
                    bool InstallPipeSignalExitHandler) {
+  if (UsageCount++)
+    return;
 #ifdef __MVS__
   // Bring stdin/stdout/stderr into a known state.
   sys::AddSignalHandler(CleanupStdHandles, nullptr);
@@ -94,6 +98,8 @@ InitLLVM::InitLLVM(int &Argc, const char **&Argv,
 }
 
 InitLLVM::~InitLLVM() {
+  if (--UsageCount)
+    return;
 #ifdef __MVS__
   CleanupStdHandles(nullptr);
 #endif
diff --git a/llvm/tools/llvm-driver/llvm-driver.cpp b/llvm/tools/llvm-driver/llvm-driver.cpp
index a0f1ca831d93b6..53a8b9357e3780 100644
--- a/llvm/tools/llvm-driver/llvm-driver.cpp
+++ b/llvm/tools/llvm-driver/llvm-driver.cpp
@@ -10,6 +10,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/LLVMDriver.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/WithColor.h"
@@ -79,4 +80,7 @@ static int findTool(int Argc, char **Argv, const char *Argv0) {
   return 1;
 }
 
-int main(int Argc, char **Argv) { return findTool(Argc, Argv, Argv[0]); }
+int main(int Argc, char **Argv) {
+  llvm::InitLLVM X(Argc, Argv);
+  return findTool(Argc, Argv, Argv[0]);
+}