4 files changed, 52 insertions, 3 deletions

diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.cpp b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
index 43aedffa808..b618e53b8e4 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.cpp
@@ -64,10 +64,12 @@ public:
   bool error(StringRef::iterator Loc, const Twine &Msg);
 
   bool parse(MachineInstr *&MI);
+  bool parseMBB(MachineBasicBlock *&MBB);
 
   bool parseRegister(unsigned &Reg);
   bool parseRegisterOperand(MachineOperand &Dest, bool IsDef = false);
   bool parseImmediateOperand(MachineOperand &Dest);
+  bool parseMBBReference(MachineBasicBlock *&MBB);
   bool parseMBBOperand(MachineOperand &Dest);
   bool parseGlobalAddressOperand(MachineOperand &Dest);
   bool parseMachineOperand(MachineOperand &Dest);
@@ -186,6 +188,19 @@ bool MIParser::parse(MachineInstr *&MI) {
   return false;
 }
 
+bool MIParser::parseMBB(MachineBasicBlock *&MBB) {
+  lex();
+  if (Token.isNot(MIToken::MachineBasicBlock))
+    return error("expected a machine basic block reference");
+  if (parseMBBReference(MBB))
+    return true;
+  lex();
+  if (Token.isNot(MIToken::Eof))
+    return error(
+        "expected end of string after the machine basic block reference");
+  return false;
+}
+
 bool MIParser::parseInstruction(unsigned &OpCode) {
   if (Token.isNot(MIToken::Identifier))
     return error("expected a machine instruction");
@@ -246,7 +261,7 @@ bool MIParser::getUnsigned(unsigned &Result) {
   return false;
 }
 
-bool MIParser::parseMBBOperand(MachineOperand &Dest) {
+bool MIParser::parseMBBReference(MachineBasicBlock *&MBB) {
   assert(Token.is(MIToken::MachineBasicBlock));
   unsigned Number;
   if (getUnsigned(Number))
@@ -255,10 +270,17 @@ bool MIParser::parseMBBOperand(MachineOperand &Dest) {
   if (MBBInfo == MBBSlots.end())
     return error(Twine("use of undefined machine basic block #") +
                  Twine(Number));
-  MachineBasicBlock *MBB = MBBInfo->second;
+  MBB = MBBInfo->second;
   if (!Token.stringValue().empty() && Token.stringValue() != MBB->getName())
     return error(Twine("the name of machine basic block #") + Twine(Number) +
                  " isn't '" + Token.stringValue() + "'");
+  return false;
+}
+
+bool MIParser::parseMBBOperand(MachineOperand &Dest) {
+  MachineBasicBlock *MBB;
+  if (parseMBBReference(MBB))
+    return true;
   Dest = MachineOperand::CreateMBB(MBB);
   lex();
   return false;
@@ -392,3 +414,10 @@ bool llvm::parseMachineInstr(
     const SlotMapping &IRSlots, SMDiagnostic &Error) {
   return MIParser(SM, MF, Error, Src, MBBSlots, IRSlots).parse(MI);
 }
+
+bool llvm::parseMBBReference(
+    MachineBasicBlock *&MBB, SourceMgr &SM, MachineFunction &MF, StringRef Src,
+    const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
+    const SlotMapping &IRSlots, SMDiagnostic &Error) {
+  return MIParser(SM, MF, Error, Src, MBBSlots, IRSlots).parseMBB(MBB);
+}
diff --git a/llvm/lib/CodeGen/MIRParser/MIParser.h b/llvm/lib/CodeGen/MIRParser/MIParser.h
index bae5c2017a2..4d6d4e70021 100644
--- a/llvm/lib/CodeGen/MIRParser/MIParser.h
+++ b/llvm/lib/CodeGen/MIRParser/MIParser.h
@@ -31,6 +31,11 @@ bool parseMachineInstr(MachineInstr *&MI, SourceMgr &SM, MachineFunction &MF,
                        const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
                        const SlotMapping &IRSlots, SMDiagnostic &Error);
 
+bool parseMBBReference(MachineBasicBlock *&MBB, SourceMgr &SM,
+                       MachineFunction &MF, StringRef Src,
+                       const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
+                       const SlotMapping &IRSlots, SMDiagnostic &Error);
+
 } // end namespace llvm
 
 #endif
diff --git a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
index 064b2f15d65..39745830078 100644
--- a/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
+++ b/llvm/lib/CodeGen/MIRParser/MIRParser.cpp
@@ -274,9 +274,18 @@ bool MIRParserImpl::initializeMachineBasicBlock(
   if (YamlMBB.AddressTaken)
     MBB.setHasAddressTaken();
   MBB.setIsLandingPad(YamlMBB.IsLandingPad);
+  SMDiagnostic Error;
+  // Parse the successors.
+  for (const auto &MBBSource : YamlMBB.Successors) {
+    MachineBasicBlock *SuccMBB = nullptr;
+    if (parseMBBReference(SuccMBB, SM, MF, MBBSource.Value, MBBSlots, IRSlots,
+                          Error))
+      return error(Error, MBBSource.SourceRange);
+    // TODO: Report an error when adding the same successor more than once.
+    MBB.addSuccessor(SuccMBB);
+  }
   // Parse the instructions.
   for (const auto &MISource : YamlMBB.Instructions) {
-    SMDiagnostic Error;
     MachineInstr *MI = nullptr;
     if (parseMachineInstr(MI, SM, MF, MISource.Value, MBBSlots, IRSlots, Error))
       return error(Error, MISource.SourceRange);
diff --git a/llvm/lib/CodeGen/MIRPrinter.cpp b/llvm/lib/CodeGen/MIRPrinter.cpp
index 18c976bcd3e..6b6675579cc 100644
--- a/llvm/lib/CodeGen/MIRPrinter.cpp
+++ b/llvm/lib/CodeGen/MIRPrinter.cpp
@@ -131,6 +131,12 @@ void MIRPrinter::convert(const Module &M, yaml::MachineBasicBlock &YamlMBB,
   YamlMBB.Alignment = MBB.getAlignment();
   YamlMBB.AddressTaken = MBB.hasAddressTaken();
   YamlMBB.IsLandingPad = MBB.isLandingPad();
+  for (const auto *MBB : MBB.successors()) {
+    std::string Str;
+    raw_string_ostream StrOS(Str);
+    MIPrinter(M, StrOS, RegisterMaskIds).printMBBReference(*MBB);
+    YamlMBB.Successors.push_back(StrOS.str());
+  }
 
   // Print the machine instructions.
   YamlMBB.Instructions.reserve(MBB.size());