[llvm-jitlink] Add Process and Platform JITDylibs, generalize alias option.

The Process JITDylib holds reflected process symbols. The Platform JITDylib
holds ORC runtime symbols if the ORC runtime is loaded. The Platform and
Process JITDylibs are appended to the link order of all other JITDylibs,
including the main JITDylib, after any explicitly specified libraries. This
scheme is similar to the one introduced in LLJIT in 371cb1af61d, and makes
it easier to introduce aliases for process and platform symbols in a way that
affects all JITDylibs uniformly.

Since the Process and Platform JITDylibs are created implicitly the -alias
option is generalized to allow source and destination JITDylibs to be explicitly
specified, i.e. the -alias option now supports general re-exports.

Testcases are updated to account for the change.

1 files changed, 157 insertions, 64 deletions

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 17fb70f3..3d6247c 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -692,11 +692,12 @@ getTestObjectFileInterface(Session &S, MemoryBufferRef O) {
 }
 
 static Error loadProcessSymbols(Session &S) {
+  S.ProcessSymsJD = &S.ES.createBareJITDylib("Process");
   auto FilterMainEntryPoint =
       [EPName = S.ES.intern(EntryPointName)](SymbolStringPtr Name) {
         return Name != EPName;
       };
-  S.MainJD->addGenerator(
+  S.ProcessSymsJD->addGenerator(
       ExitOnErr(orc::EPCDynamicLibrarySearchGenerator::GetForTargetProcess(
           S.ES, std::move(FilterMainEntryPoint))));
 
@@ -707,8 +708,9 @@ static Error loadDylibs(Session &S) {
   LLVM_DEBUG(dbgs() << "Loading dylibs...\n");
   for (const auto &Dylib : Dylibs) {
     LLVM_DEBUG(dbgs() << "  " << Dylib << "\n");
-    if (auto Err = S.loadAndLinkDynamicLibrary(*S.MainJD, Dylib))
-      return Err;
+    auto DL = S.getOrLoadDynamicLibrary(Dylib);
+    if (!DL)
+      return DL.takeError();
   }
 
   return Error::success();
@@ -963,69 +965,79 @@ Session::Session(std::unique_ptr<ExecutorProcessControl> EPC, Error &Err)
 
   ES.setErrorReporter(reportLLVMJITLinkError);
 
-  if (auto MainJDOrErr = ES.createJITDylib("main"))
-    MainJD = &*MainJDOrErr;
-  else {
-    Err = MainJDOrErr.takeError();
-    return;
-  }
-
   if (!NoProcessSymbols)
     ExitOnErr(loadProcessSymbols(*this));
-  else {
-    // This symbol is used in testcases.
-    auto &TestResultJD = ES.createBareJITDylib("<TestResultJD>");
-    ExitOnErr(TestResultJD.define(absoluteSymbols(
-        {{ES.intern("llvm_jitlink_setTestResultOverride"),
-          {ExecutorAddr::fromPtr(llvm_jitlink_setTestResultOverride),
-           JITSymbolFlags::Exported}}})));
-    MainJD->addToLinkOrder(TestResultJD);
-  }
 
   ExitOnErr(loadDylibs(*this));
 
   auto &TT = ES.getTargetTriple();
 
-  if (DebuggerSupport && TT.isOSBinFormatMachO())
-    ObjLayer.addPlugin(ExitOnErr(
-        GDBJITDebugInfoRegistrationPlugin::Create(this->ES, *MainJD, TT)));
+  if (DebuggerSupport && TT.isOSBinFormatMachO()) {
+    if (!ProcessSymsJD) {
+      Err = make_error<StringError>("MachO debugging requires process symbols",
+                                    inconvertibleErrorCode());
+      return;
+    }
+    ObjLayer.addPlugin(ExitOnErr(GDBJITDebugInfoRegistrationPlugin::Create(
+        this->ES, *ProcessSymsJD, TT)));
+  }
 
   if (PerfSupport && TT.isOSBinFormatELF()) {
+    if (!ProcessSymsJD) {
+      Err = make_error<StringError>("MachO debugging requires process symbols",
+                                    inconvertibleErrorCode());
+      return;
+    }
     ObjLayer.addPlugin(ExitOnErr(DebugInfoPreservationPlugin::Create()));
     ObjLayer.addPlugin(ExitOnErr(PerfSupportPlugin::Create(
-        this->ES.getExecutorProcessControl(), *MainJD, true, true)));
+        this->ES.getExecutorProcessControl(), *ProcessSymsJD, true, true)));
   }
 
   // Set up the platform.
-  if (TT.isOSBinFormatMachO() && !OrcRuntime.empty()) {
-    if (auto P =
-            MachOPlatform::Create(ES, ObjLayer, *MainJD, OrcRuntime.c_str()))
-      ES.setPlatform(std::move(*P));
-    else {
-      Err = P.takeError();
-      return;
-    }
-  } else if (TT.isOSBinFormatELF() && !OrcRuntime.empty()) {
-    if (auto P =
-            ELFNixPlatform::Create(ES, ObjLayer, *MainJD, OrcRuntime.c_str()))
-      ES.setPlatform(std::move(*P));
-    else {
-      Err = P.takeError();
-      return;
-    }
-  } else if (TT.isOSBinFormatCOFF() && !OrcRuntime.empty()) {
-    auto LoadDynLibrary = [&, this](JITDylib &JD, StringRef DLLName) -> Error {
-      if (!DLLName.ends_with_insensitive(".dll"))
-        return make_error<StringError>("DLLName not ending with .dll",
-                                       inconvertibleErrorCode());
-      return loadAndLinkDynamicLibrary(JD, DLLName);
-    };
+  if (!OrcRuntime.empty()) {
+    assert(ProcessSymsJD && "ProcessSymsJD should have been set");
+    PlatformJD = &ES.createBareJITDylib("Platform");
+    PlatformJD->addToLinkOrder(*ProcessSymsJD);
+
+    if (TT.isOSBinFormatMachO()) {
+      if (auto P = MachOPlatform::Create(ES, ObjLayer, *PlatformJD,
+                                         OrcRuntime.c_str()))
+        ES.setPlatform(std::move(*P));
+      else {
+        Err = P.takeError();
+        return;
+      }
+    } else if (TT.isOSBinFormatELF()) {
+      if (auto P = ELFNixPlatform::Create(ES, ObjLayer, *PlatformJD,
+                                          OrcRuntime.c_str()))
+        ES.setPlatform(std::move(*P));
+      else {
+        Err = P.takeError();
+        return;
+      }
+    } else if (TT.isOSBinFormatCOFF()) {
+      auto LoadDynLibrary = [&, this](JITDylib &JD,
+                                      StringRef DLLName) -> Error {
+        if (!DLLName.ends_with_insensitive(".dll"))
+          return make_error<StringError>("DLLName not ending with .dll",
+                                         inconvertibleErrorCode());
+        return loadAndLinkDynamicLibrary(JD, DLLName);
+      };
 
-    if (auto P = COFFPlatform::Create(ES, ObjLayer, *MainJD, OrcRuntime.c_str(),
-                                      std::move(LoadDynLibrary))) 
-      ES.setPlatform(std::move(*P));
-     else {
-      Err = P.takeError();
+      if (auto P = COFFPlatform::Create(ES, ObjLayer, *PlatformJD,
+                                        OrcRuntime.c_str(),
+                                        std::move(LoadDynLibrary)))
+        ES.setPlatform(std::move(*P));
+      else {
+        Err = P.takeError();
+        return;
+      }
+    } else {
+      Err = make_error<StringError>(
+          "-" + OrcRuntime.ArgStr + " specified, but format " +
+              Triple::getObjectFormatTypeName(TT.getObjectFormat()) +
+              " not supported",
+          inconvertibleErrorCode());
       return;
     }
   } else if (TT.isOSBinFormatELF()) {
@@ -1037,6 +1049,24 @@ Session::Session(std::unique_ptr<ExecutorProcessControl> EPC, Error &Err)
           ES, ExitOnErr(createJITLoaderGDBRegistrar(this->ES)), true, true));
   }
 
+  if (auto MainJDOrErr = ES.createJITDylib("main"))
+    MainJD = &*MainJDOrErr;
+  else {
+    Err = MainJDOrErr.takeError();
+    return;
+  }
+
+  if (NoProcessSymbols) {
+    // This symbol is used in testcases, but we're not reflecting process
+    // symbols so we'll need to make it available some other way.
+    auto &TestResultJD = ES.createBareJITDylib("<TestResultJD>");
+    ExitOnErr(TestResultJD.define(absoluteSymbols(
+        {{ES.intern("llvm_jitlink_setTestResultOverride"),
+          {ExecutorAddr::fromPtr(llvm_jitlink_setTestResultOverride),
+           JITSymbolFlags::Exported}}})));
+    MainJD->addToLinkOrder(TestResultJD);
+  }
+
   ObjLayer.addPlugin(std::make_unique<JITLinkSessionPlugin>(*this));
 
   // Process any harness files.
@@ -1266,6 +1296,10 @@ static Error sanitizeArguments(const Triple &TT, const char *ArgV0) {
   if (DebuggerSupport.getNumOccurrences() == 0 && NoExec)
     DebuggerSupport = false;
 
+  if (!OrcRuntime.empty() && NoProcessSymbols)
+    return make_error<StringError>("-orc-runtime requires process symbols",
+                                   inconvertibleErrorCode());
+
   // If -slab-allocate is passed, check that we're not trying to use it in
   // -oop-executor or -oop-executor-connect mode.
   //
@@ -1365,6 +1399,13 @@ static Error createJITDylibs(Session &S,
     }
   }
 
+  if (S.PlatformJD)
+    S.JDSearchOrder.push_back(
+        {S.PlatformJD, JITDylibLookupFlags::MatchExportedSymbolsOnly});
+  if (S.ProcessSymsJD)
+    S.JDSearchOrder.push_back(
+        {S.ProcessSymsJD, JITDylibLookupFlags::MatchExportedSymbolsOnly});
+
   LLVM_DEBUG({
     dbgs() << "Dylib search order is [ ";
     for (auto &KV : S.JDSearchOrder)
@@ -1416,23 +1457,67 @@ static Error addAliases(Session &S,
                         const std::map<unsigned, JITDylib *> &IdxToJD) {
   // Define absolute symbols.
   LLVM_DEBUG(dbgs() << "Defining aliases...\n");
+
+  DenseMap<std::pair<JITDylib *, JITDylib *>, SymbolAliasMap> Reexports;
   for (auto AliasItr = Aliases.begin(), AliasEnd = Aliases.end();
        AliasItr != AliasEnd; ++AliasItr) {
-    unsigned AliasArgIdx = Aliases.getPosition(AliasItr - Aliases.begin());
-    auto &JD = *std::prev(IdxToJD.lower_bound(AliasArgIdx))->second;
 
-    StringRef AliasStmt = *AliasItr;
-    size_t EqIdx = AliasStmt.find_first_of('=');
-    if (EqIdx == StringRef::npos)
-      return make_error<StringError>("Invalid alias definition \"" + AliasStmt +
-                                         "\". Syntax: <name>=<addr>",
-                                     inconvertibleErrorCode());
-    StringRef Alias = AliasStmt.substr(0, EqIdx).trim();
-    StringRef Aliasee = AliasStmt.substr(EqIdx + 1).trim();
+    auto BadExpr = [&]() {
+      return make_error<StringError>(
+          "Invalid alias definition \"" + *AliasItr +
+              "\". Syntax: [<dst-jd>:]<alias>=[<src-jd>:]<aliasee>",
+          inconvertibleErrorCode());
+    };
+
+    auto GetJD = [&](StringRef JDName) -> Expected<JITDylib *> {
+      if (JDName.empty()) {
+        unsigned AliasArgIdx = Aliases.getPosition(AliasItr - Aliases.begin());
+        return std::prev(IdxToJD.lower_bound(AliasArgIdx))->second;
+      }
+
+      auto *JD = S.ES.getJITDylibByName(JDName);
+      if (!JD)
+        return make_error<StringError>(StringRef("In alias definition \"") +
+                                           *AliasItr + "\" no dylib named " +
+                                           JDName,
+                                       inconvertibleErrorCode());
 
-    SymbolAliasMap SAM;
-    SAM[S.ES.intern(Alias)] = {S.ES.intern(Aliasee), JITSymbolFlags::Exported};
-    if (auto Err = JD.define(symbolAliases(std::move(SAM))))
+      return JD;
+    };
+
+    {
+      // First split on '=' to get alias and aliasee.
+      StringRef AliasStmt = *AliasItr;
+      auto [AliasExpr, AliaseeExpr] = AliasStmt.split('=');
+      if (AliaseeExpr.empty())
+        return BadExpr();
+
+      auto [AliasJDName, Alias] = AliasExpr.split(':');
+      if (Alias.empty())
+        std::swap(AliasJDName, Alias);
+
+      auto AliasJD = GetJD(AliasJDName);
+      if (!AliasJD)
+        return AliasJD.takeError();
+
+      auto [AliaseeJDName, Aliasee] = AliaseeExpr.split(':');
+      if (Aliasee.empty())
+        std::swap(AliaseeJDName, Aliasee);
+
+      if (AliaseeJDName.empty() && !AliasJDName.empty())
+        AliaseeJDName = AliasJDName;
+      auto AliaseeJD = GetJD(AliaseeJDName);
+      if (!AliaseeJD)
+        return AliaseeJD.takeError();
+
+      Reexports[{*AliasJD, *AliaseeJD}][S.ES.intern(Alias)] = {
+          S.ES.intern(Aliasee), JITSymbolFlags::Exported};
+    }
+  }
+
+  for (auto &[JDs, AliasMap] : Reexports) {
+    auto [DstJD, SrcJD] = JDs;
+    if (auto Err = DstJD->define(reexports(*SrcJD, std::move(AliasMap))))
       return Err;
   }
 
@@ -1766,6 +1851,14 @@ static Error addLibraries(Session &S,
                                      inconvertibleErrorCode());
   }
 
+  // Add platform and process symbols if available.
+  for (auto &[Idx, JD] : IdxToJD) {
+    if (S.PlatformJD)
+      JD->addToLinkOrder(*S.PlatformJD);
+    if (S.ProcessSymsJD)
+      JD->addToLinkOrder(*S.ProcessSymsJD);
+  }
+
   return Error::success();
 }