[IR] Support importing modules with invalid data layouts.

Use the existing mechanism to change the data layout using callbacks.

Before this patch, we had a callback type DataLayoutCallbackTy that receives
a single StringRef specifying the target triple, and optionally returns
the data layout string to be used. Module loaders (both IR and BC) then
apply the callback to potentially override the module's data layout,
after first having imported and parsed the data layout from the file.

We can't do the same to fix invalid data layouts, because the import will already
fail, before the callback has a chance to fix it.
Instead, module loaders now tentatively parse the data layout into a string,
wait until the target triple has been parsed, apply the override callback
to the imported string and only then parse the tentative string as a data layout.

Moreover, add the old data layout string S as second argument to the callback,
in addition to the already existing target triple argument.
S is either the default data layout string in case none is specified, or the data
layout string specified in the module, possibly after auto-upgrades (for the BitcodeReader).
This allows callbacks to inspect the old data layout string,
and fix it instead of setting a fixed data layout.

Also allow to pass data layout override callbacks to lazy bitcode module
loader functions.

Differential Revision: https://reviews.llvm.org/D140985

Author: jasilvanus

llvm/include/llvm/AsmParser/LLParser.h

@@ -179,8 +179,10 @@ namespace llvm {
           Lex(F, SM, Err, Context), M(M), Index(Index), Slots(Slots),
           BlockAddressPFS(nullptr) {}
     bool Run(
-        bool UpgradeDebugInfo, DataLayoutCallbackTy DataLayoutCallback =
-                                   [](StringRef) { return std::nullopt; });
+        bool UpgradeDebugInfo,
+        DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+          return std::nullopt;
+        });
 
     bool parseStandaloneConstantValue(Constant *&C, const SlotMapping *Slots);
 
@@ -318,8 +320,8 @@ namespace llvm {
     bool parseTopLevelEntities();
     bool validateEndOfModule(bool UpgradeDebugInfo);
     bool validateEndOfIndex();
-    bool parseTargetDefinitions();
-    bool parseTargetDefinition();
+    bool parseTargetDefinitions(DataLayoutCallbackTy DataLayoutCallback);
+    bool parseTargetDefinition(std::string &TentativeDLStr, LocTy &DLStrLoc);
     bool parseModuleAsm();
     bool parseSourceFileName();
     bool parseUnnamedType();

llvm/include/llvm/AsmParser/Parser.h

@@ -29,7 +29,7 @@ struct SlotMapping;
 class SMDiagnostic;
 class Type;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
 /// This function is a main interface to the LLVM Assembly Parser. It parses
@@ -86,7 +86,7 @@ struct ParsedModuleAndIndex {
 ParsedModuleAndIndex parseAssemblyFileWithIndex(
     StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 
@@ -127,7 +127,7 @@ parseSummaryIndexAssemblyString(StringRef AsmString, SMDiagnostic &Err);
 std::unique_ptr<Module> parseAssembly(
     MemoryBufferRef F, SMDiagnostic &Err, LLVMContext &Context,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 
@@ -169,7 +169,7 @@ parseSummaryIndexAssembly(MemoryBufferRef F, SMDiagnostic &Err);
 bool parseAssemblyInto(
     MemoryBufferRef F, Module *M, ModuleSummaryIndex *Index, SMDiagnostic &Err,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 

llvm/include/llvm/Bitcode/BitcodeReader.h

@@ -34,7 +34,11 @@ class Module;
 class MemoryBuffer;
 class ModuleSummaryIndex;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+// Callback to override the data layout string of an imported bitcode module.
+// The first argument is the target triple, the second argument the data layout
+// string from the input, or a default string. It will be used if the callback
+// returns std::nullopt.
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
   // These functions are for converting Expected/Error values to
@@ -101,14 +105,18 @@ typedef llvm::function_ref<std::optional<std::string>(StringRef)>
     /// bodies. If ShouldLazyLoadMetadata is true, lazily load metadata as well.
     /// If IsImporting is true, this module is being parsed for ThinLTO
     /// importing into another module.
-    Expected<std::unique_ptr<Module>> getLazyModule(LLVMContext &Context,
-                                                    bool ShouldLazyLoadMetadata,
-                                                    bool IsImporting);
+    Expected<std::unique_ptr<Module>> getLazyModule(
+        LLVMContext &Context, bool ShouldLazyLoadMetadata, bool IsImporting,
+        DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+          return std::nullopt;
+        });
 
     /// Read the entire bitcode module and return it.
     Expected<std::unique_ptr<Module>> parseModule(
-        LLVMContext &Context, DataLayoutCallbackTy DataLayoutCallback =
-                                  [](StringRef) { return std::nullopt; });
+        LLVMContext &Context,
+        DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+          return std::nullopt;
+        });
 
     /// Returns information about the module to be used for LTO: whether to
     /// compile with ThinLTO, and whether it has a summary.
@@ -145,10 +153,12 @@ typedef llvm::function_ref<std::optional<std::string>(StringRef)>
   /// deserialization of function bodies. If ShouldLazyLoadMetadata is true,
   /// lazily load metadata as well. If IsImporting is true, this module is
   /// being parsed for ThinLTO importing into another module.
-  Expected<std::unique_ptr<Module>>
-  getLazyBitcodeModule(MemoryBufferRef Buffer, LLVMContext &Context,
-                       bool ShouldLazyLoadMetadata = false,
-                       bool IsImporting = false);
+  Expected<std::unique_ptr<Module>> getLazyBitcodeModule(
+      MemoryBufferRef Buffer, LLVMContext &Context,
+      bool ShouldLazyLoadMetadata = false, bool IsImporting = false,
+      DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
+        return std::nullopt;
+      });
 
   /// Like getLazyBitcodeModule, except that the module takes ownership of
   /// the memory buffer if successful. If successful, this moves Buffer. On
@@ -175,7 +185,7 @@ typedef llvm::function_ref<std::optional<std::string>(StringRef)>
   /// Read the specified bitcode file, returning the module.
   Expected<std::unique_ptr<Module>> parseBitcodeFile(
       MemoryBufferRef Buffer, LLVMContext &Context,
-      DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+      DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
         return std::nullopt;
       });
 

llvm/include/llvm/CodeGen/MIRParser/MIRParser.h

@@ -34,7 +34,7 @@ class MachineModuleInfo;
 class SMDiagnostic;
 class StringRef;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
 /// This class initializes machine functions by applying the state loaded from
@@ -52,9 +52,8 @@ class MIRParser {
   /// A new, empty module is created if the LLVM IR isn't present.
   /// \returns nullptr if a parsing error occurred.
   std::unique_ptr<Module>
-  parseIRModule(DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
-    return std::nullopt;
-  });
+  parseIRModule(DataLayoutCallbackTy DataLayoutCallback =
+                    [](StringRef, StringRef) { return std::nullopt; });
 
   /// Parses MachineFunctions in the MIR file and add them to the given
   /// MachineModuleInfo \p MMI.

llvm/include/llvm/IRReader/IRReader.h

@@ -27,7 +27,7 @@ class Module;
 class SMDiagnostic;
 class LLVMContext;
 
-typedef llvm::function_ref<std::optional<std::string>(StringRef)>
+typedef llvm::function_ref<std::optional<std::string>(StringRef, StringRef)>
     DataLayoutCallbackTy;
 
 /// If the given MemoryBuffer holds a bitcode image, return a Module
@@ -55,7 +55,7 @@ getLazyIRFileModule(StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
 std::unique_ptr<Module> parseIR(
     MemoryBufferRef Buffer, SMDiagnostic &Err, LLVMContext &Context,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 
@@ -65,7 +65,7 @@ std::unique_ptr<Module> parseIR(
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
 std::unique_ptr<Module> parseIRFile(
     StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef) {
+    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
       return std::nullopt;
     });
 }

llvm/lib/AsmParser/LLParser.cpp

@@ -95,11 +95,8 @@ bool LLParser::Run(bool UpgradeDebugInfo,
         "Can't read textual IR with a Context that discards named Values");
 
   if (M) {
-    if (parseTargetDefinitions())
+    if (parseTargetDefinitions(DataLayoutCallback))
       return true;
-
-    if (auto LayoutOverride = DataLayoutCallback(M->getTargetTriple()))
-      M->setDataLayout(*LayoutOverride);
   }
 
   return parseTopLevelEntities() || validateEndOfModule(UpgradeDebugInfo) ||
@@ -353,21 +350,41 @@ bool LLParser::validateEndOfIndex() {
 // Top-Level Entities
 //===----------------------------------------------------------------------===//
 
-bool LLParser::parseTargetDefinitions() {
-  while (true) {
+bool LLParser::parseTargetDefinitions(DataLayoutCallbackTy DataLayoutCallback) {
+  // Delay parsing of the data layout string until the target triple is known.
+  // Then, pass both the the target triple and the tentative data layout string
+  // to DataLayoutCallback, allowing to override the DL string.
+  // This enables importing modules with invalid DL strings.
+  std::string TentativeDLStr = M->getDataLayoutStr();
+  LocTy DLStrLoc;
+
+  bool Done = false;
+  while (!Done) {
     switch (Lex.getKind()) {
     case lltok::kw_target:
-      if (parseTargetDefinition())
+      if (parseTargetDefinition(TentativeDLStr, DLStrLoc))
         return true;
       break;
     case lltok::kw_source_filename:
       if (parseSourceFileName())
         return true;
       break;
     default:
-      return false;
+      Done = true;
     }
   }
+  // Run the override callback to potentially change the data layout string, and
+  // parse the data layout string.
+  if (auto LayoutOverride =
+          DataLayoutCallback(M->getTargetTriple(), TentativeDLStr)) {
+    TentativeDLStr = *LayoutOverride;
+    DLStrLoc = {};
+  }
+  Expected<DataLayout> MaybeDL = DataLayout::parse(TentativeDLStr);
+  if (!MaybeDL)
+    return error(DLStrLoc, toString(MaybeDL.takeError()));
+  M->setDataLayout(MaybeDL.get());
+  return false;
 }
 
 bool LLParser::parseTopLevelEntities() {
@@ -471,7 +488,8 @@ bool LLParser::parseModuleAsm() {
 /// toplevelentity
 ///   ::= 'target' 'triple' '=' STRINGCONSTANT
 ///   ::= 'target' 'datalayout' '=' STRINGCONSTANT
-bool LLParser::parseTargetDefinition() {
+bool LLParser::parseTargetDefinition(std::string &TentativeDLStr,
+                                     LocTy &DLStrLoc) {
   assert(Lex.getKind() == lltok::kw_target);
   std::string Str;
   switch (Lex.Lex()) {
@@ -488,13 +506,9 @@ bool LLParser::parseTargetDefinition() {
     Lex.Lex();
     if (parseToken(lltok::equal, "expected '=' after target datalayout"))
       return true;
-    LocTy Loc = Lex.getLoc();
-    if (parseStringConstant(Str))
+    DLStrLoc = Lex.getLoc();
+    if (parseStringConstant(TentativeDLStr))
       return true;
-    Expected<DataLayout> MaybeDL = DataLayout::parse(Str);
-    if (!MaybeDL)
-      return error(Loc, toString(MaybeDL.takeError()));
-    M->setDataLayout(MaybeDL.get());
     return false;
   }
 }

llvm/lib/AsmParser/Parser.cpp

@@ -91,9 +91,10 @@ ParsedModuleAndIndex llvm::parseAssemblyWithIndex(MemoryBufferRef F,
                                                   SMDiagnostic &Err,
                                                   LLVMContext &Context,
                                                   SlotMapping *Slots) {
-  return ::parseAssemblyWithIndex(F, Err, Context, Slots,
-                                  /*UpgradeDebugInfo*/ true,
-                                  [](StringRef) { return std::nullopt; });
+  return ::parseAssemblyWithIndex(
+      F, Err, Context, Slots,
+      /*UpgradeDebugInfo*/ true,
+      [](StringRef, StringRef) { return std::nullopt; });
 }
 
 static ParsedModuleAndIndex
@@ -150,7 +151,7 @@ static bool parseSummaryIndexAssemblyInto(MemoryBufferRef F,
   // index, but we need to initialize it.
   LLVMContext unusedContext;
   return LLParser(F.getBuffer(), SM, Err, nullptr, &Index, unusedContext)
-      .Run(true, [](StringRef) { return std::nullopt; });
+      .Run(true, [](StringRef, StringRef) { return std::nullopt; });
 }
 
 std::unique_ptr<ModuleSummaryIndex>