1 files changed, 58 insertions, 0 deletions

diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
index c46adf5c719..1bdaedb5d3a 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -23,6 +23,7 @@
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Module.h"
+#include "llvm/IR/ModuleSlotTracker.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Target/TargetSubtargetInfo.h"
@@ -77,6 +78,8 @@ class MIParser {
   StringMap<const uint32_t *> Names2RegMasks;
   /// Maps from subregister names to subregister indices.
   StringMap<unsigned> Names2SubRegIndices;
+  /// Maps from slot numbers to function's unnamed basic blocks.
+  DenseMap<unsigned, const BasicBlock *> Slots2BasicBlocks;
 
 public:
   MIParser(SourceMgr &SM, MachineFunction &MF, SMDiagnostic &Error,
@@ -99,6 +102,7 @@ public:
   bool parseStandaloneMBB(MachineBasicBlock *&MBB);
   bool parseStandaloneNamedRegister(unsigned &Reg);
   bool parseStandaloneVirtualRegister(unsigned &Reg);
+  bool parseStandaloneIRBlockReference(const BasicBlock *&BB);
 
   bool parseRegister(unsigned &Reg);
   bool parseRegisterFlag(unsigned &Flags);
@@ -160,6 +164,10 @@ private:
   ///
   /// Return 0 if the name isn't a subregister index class.
   unsigned getSubRegIndex(StringRef Name);
+
+  void initSlots2BasicBlocks();
+
+  const BasicBlock *getIRBlock(unsigned Slot);
 };
 
 } // end anonymous namespace
@@ -302,6 +310,23 @@ bool MIParser::parseStandaloneVirtualRegister(unsigned &Reg) {
   return false;
 }
 
+bool MIParser::parseStandaloneIRBlockReference(const BasicBlock *&BB) {
+  lex();
+  if (Token.isNot(MIToken::IRBlock))
+    return error("expected an IR block reference");
+  unsigned SlotNumber = 0;
+  if (getUnsigned(SlotNumber))
+    return true;
+  BB = getIRBlock(SlotNumber);
+  if (!BB)
+    return error(Twine("use of undefined IR block '%ir-block.") +
+                 Twine(SlotNumber) + "'");
+  lex();
+  if (Token.isNot(MIToken::Eof))
+    return error("expected end of string after the IR block reference");
+  return false;
+}
+
 static const char *printImplicitRegisterFlag(const MachineOperand &MO) {
   assert(MO.isImplicit());
   return MO.isDef() ? "implicit-def" : "implicit";
@@ -842,6 +867,30 @@ unsigned MIParser::getSubRegIndex(StringRef Name) {
   return SubRegInfo->getValue();
 }
 
+void MIParser::initSlots2BasicBlocks() {
+  if (!Slots2BasicBlocks.empty())
+    return;
+  const auto &F = *MF.getFunction();
+  ModuleSlotTracker MST(F.getParent());
+  MST.incorporateFunction(F);
+  for (auto &BB : F) {
+    if (BB.hasName())
+      continue;
+    int Slot = MST.getLocalSlot(&BB);
+    if (Slot == -1)
+      continue;
+    Slots2BasicBlocks.insert(std::make_pair(unsigned(Slot), &BB));
+  }
+}
+
+const BasicBlock *MIParser::getIRBlock(unsigned Slot) {
+  initSlots2BasicBlocks();
+  auto BlockInfo = Slots2BasicBlocks.find(Slot);
+  if (BlockInfo == Slots2BasicBlocks.end())
+    return nullptr;
+  return BlockInfo->second;
+}
+
 bool llvm::parseMachineInstr(MachineInstr *&MI, SourceMgr &SM,
                              MachineFunction &MF, StringRef Src,
                              const PerFunctionMIParsingState &PFS,
@@ -873,3 +922,12 @@ bool llvm::parseVirtualRegisterReference(unsigned &Reg, SourceMgr &SM,
   return MIParser(SM, MF, Error, Src, PFS, IRSlots)
       .parseStandaloneVirtualRegister(Reg);
 }
+
+bool llvm::parseIRBlockReference(const BasicBlock *&BB, SourceMgr &SM,
+                                 MachineFunction &MF, StringRef Src,
+                                 const PerFunctionMIParsingState &PFS,
+                                 const SlotMapping &IRSlots,
+                                 SMDiagnostic &Error) {
+  return MIParser(SM, MF, Error, Src, PFS, IRSlots)
+      .parseStandaloneIRBlockReference(BB);
+}