[ORC] Add a runAsMain utility function to ExecutionUtils.

The runAsMain function takes a pointer to a function with a standard C main
signature, int(*)(int, char*[]), and invokes it using the given arguments and
program name. The arguments are copied into writable temporary storage as
required by the C and C++ specifications, so runAsMain safe to use when calling
main functions that modify their arguments in-place.

This patch also uses the new runAsMain function to replace hand-rolled versions
in lli, llvm-jitlink, and the SpeculativeJIT example.

5 files changed, 48 insertions, 42 deletions

diff --git a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp
index 1fd1fc92a73..5a0ad376bcd 100644
--- a/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp
+++ b/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp
@@ -176,20 +176,11 @@ int main(int argc, char *argv[]) {
                             ThreadSafeModule(std::move(M), std::move(Ctx))));
   }
 
-  // Build an argv array for the JIT'd main.
-  std::vector<const char *> ArgV;
-  ArgV.push_back(argv[0]);
-  for (const auto &InputArg : InputArgv)
-    ArgV.push_back(InputArg.data());
-  ArgV.push_back(nullptr);
-
-  // Look up the JIT'd main, cast it to a function pointer, then call it.
-
   auto MainSym = ExitOnErr(SJ->lookup("main"));
-  int (*Main)(int, const char *[]) =
-      (int (*)(int, const char *[]))MainSym.getAddress();
+  auto Main =
+      jitTargetAddressToFunction<int (*)(int, char *[])>(MainSym.getAddress());
 
-  Main(ArgV.size() - 1, ArgV.data());
+  return runAsMain(Main, InputArgv, StringRef(InputFiles.front()));
 
   return 0;
 }
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
index 7d051ed9906..c797dbbbdfd 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ExecutionUtils.h
@@ -40,6 +40,17 @@ namespace orc {
 
 class ObjectLayer;
 
+/// Run a main function, returning the result.
+///
+/// If the optional ProgramName argument is given then it will be inserted
+/// before the strings in Args as the first argument to the called function.
+///
+/// It is legal to have an empty argument list and no program name, however
+/// many main functions will expect a name argument at least, and will fail
+/// if none is provided.
+int runAsMain(int (*Main)(int, char *[]), ArrayRef<std::string> Args,
+              Optional<StringRef> ProgramName = None);
+
 /// This iterator provides a convenient way to iterate over the elements
 ///        of an llvm.global_ctors/llvm.global_dtors instance.
 ///
diff --git a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
index 4a3482242db..0a3fef207ac 100644
--- a/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ExecutionUtils.cpp
@@ -19,6 +19,32 @@
 namespace llvm {
 namespace orc {
 
+int runAsMain(int (*Main)(int, char *[]), ArrayRef<std::string> Args,
+              Optional<StringRef> ProgramName) {
+  std::vector<std::unique_ptr<char[]>> ArgVStorage;
+  std::vector<char *> ArgV;
+
+  ArgVStorage.reserve(Args.size() + (ProgramName ? 1 : 0));
+  ArgV.reserve(Args.size() + 1 + (ProgramName ? 1 : 0));
+
+  if (ProgramName) {
+    ArgVStorage.push_back(std::make_unique<char[]>(ProgramName->size() + 1));
+    llvm::copy(*ProgramName, &ArgVStorage.back()[0]);
+    ArgVStorage.back()[ProgramName->size()] = '\0';
+    ArgV.push_back(ArgVStorage.back().get());
+  }
+
+  for (auto &Arg : Args) {
+    ArgVStorage.push_back(std::make_unique<char[]>(Arg.size() + 1));
+    llvm::copy(Arg, &ArgVStorage.back()[0]);
+    ArgVStorage.back()[Arg.size()] = '\0';
+    ArgV.push_back(ArgVStorage.back().get());
+  }
+  ArgV.push_back(nullptr);
+
+  return Main(Args.size(), ArgV.data());
+}
+
 CtorDtorIterator::CtorDtorIterator(const GlobalVariable *GV, bool End)
   : InitList(
       GV ? dyn_cast_or_null<ConstantArray>(GV->getInitializer()) : nullptr),
diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp
index ccad0672141..9aaeef08592 100644
--- a/llvm/tools/lli/lli.cpp
+++ b/llvm/tools/lli/lli.cpp
@@ -871,16 +871,11 @@ int runOrcLazyJIT(const char *ProgName) {
 
   // Run main.
   auto MainSym = ExitOnErr(J->lookup("main"));
-  typedef int (*MainFnPtr)(int, const char *[]);
-  std::vector<const char *> ArgV;
-  for (auto &Arg : Args)
-    ArgV.push_back(Arg.c_str());
-  ArgV.push_back(nullptr);
-
-  int ArgC = ArgV.size() - 1;
-  auto Main =
-      reinterpret_cast<MainFnPtr>(static_cast<uintptr_t>(MainSym.getAddress()));
-  auto Result = Main(ArgC, (const char **)ArgV.data());
+
+  typedef int (*MainFnPtr)(int, char *[]);
+  auto Result = orc::runAsMain(
+      jitTargetAddressToFunction<MainFnPtr>(MainSym.getAddress()), Args,
+      StringRef("lli"));
 
   // Wait for -entry-point threads.
   for (auto &AltEntryThread : AltEntryThreads)
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 003a333d956..0e0953d37b6 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -769,25 +769,6 @@ static Expected<JITEvaluatedSymbol> getMainEntryPoint(Session &S) {
   return S.ES.lookup(S.JDSearchOrder, EntryPointName);
 }
 
-Expected<int> runEntryPoint(Session &S, JITEvaluatedSymbol EntryPoint) {
-  assert(EntryPoint.getAddress() && "Entry point address should not be null");
-
-  constexpr const char *JITProgramName = "<llvm-jitlink jit'd code>";
-  auto PNStorage = std::make_unique<char[]>(strlen(JITProgramName) + 1);
-  strcpy(PNStorage.get(), JITProgramName);
-
-  std::vector<const char *> EntryPointArgs;
-  EntryPointArgs.push_back(PNStorage.get());
-  for (auto &InputArg : InputArgv)
-    EntryPointArgs.push_back(InputArg.data());
-  EntryPointArgs.push_back(nullptr);
-
-  using MainTy = int (*)(int, const char *[]);
-  MainTy EntryPointPtr = reinterpret_cast<MainTy>(EntryPoint.getAddress());
-
-  return EntryPointPtr(EntryPointArgs.size() - 1, EntryPointArgs.data());
-}
-
 struct JITLinkTimers {
   TimerGroup JITLinkTG{"llvm-jitlink timers", "timers for llvm-jitlink phases"};
   Timer LoadObjectsTimer{"load", "time to load/add object files", JITLinkTG};
@@ -841,8 +822,10 @@ int main(int argc, char *argv[]) {
 
   int Result = 0;
   {
+    using MainTy = int (*)(int, char *[]);
+    auto EntryFn = jitTargetAddressToFunction<MainTy>(EntryPoint.getAddress());
     TimeRegion TR(Timers ? &Timers->RunTimer : nullptr);
-    Result = ExitOnErr(runEntryPoint(S, EntryPoint));
+    Result = runAsMain(EntryFn, InputArgv, StringRef(InputFiles.front()));
   }
 
   return Result;