MIR Serialization: Serialize MBB Successors.

[arphaman]

Summary

This patch serializes the machine basic block successors. It uses a YAML flow sequence that contains
strings that have the MBB references. The MBB reference strings use the same syntax as MBB machine
operands in the machine instruction strings.

Example:

body:
  - id:          0
    name:    entry
    successors:  [ '%bb.1.less', '%bb.2.exit' ]
  - id:          1
    name:    less
  - id:          2
    name:    exit

[arphaman]

I rebased this patch on top of trunk.

[arphaman]

Cleanup this patch taking Duncan's comments into account.

include/llvm/CodeGen/MIRYamlMapping.h

if (const auto *Node =
             reinterpret_cast<yaml::Input *>(Ctx)->getCurrentNode())
       S.SourceRange = Node->getSourceRange();
     return "";
   }
 
   static bool mustQuote(StringRef Scalar) { return needsQuotes(Scalar); }
 };
 
+struct FlowStringValue : StringValue {
+  FlowStringValue() {}
+  FlowStringValue(std::string Value) : StringValue(Value) {}
+};
+
+template <> struct ScalarTraits<FlowStringValue> {
+  static void output(const FlowStringValue &S, void *, llvm::raw_ostream &OS) {
+    return ScalarTraits<StringValue>::output(S, nullptr, OS);
+  }
+
+  static StringRef input(StringRef Scalar, void *Ctx, FlowStringValue &S) {
+    return ScalarTraits<StringValue>::input(Scalar, Ctx, S);
+  }
+
+  static bool mustQuote(StringRef Scalar) { return needsQuotes(Scalar); }
+};
+
 } // end namespace yaml
 } // end namespace llvm
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::StringValue)
+LLVM_YAML_IS_FLOW_SEQUENCE_VECTOR(llvm::yaml::FlowStringValue)
 
 namespace llvm {
 namespace yaml {
 
 struct MachineBasicBlock {
   unsigned ID;
   std::string Name;
   unsigned Alignment = 0;
   bool IsLandingPad = false;
   bool AddressTaken = false;
-  // TODO: Serialize the successors and liveins.
+  // TODO: Serialize the successor weights and liveins.
+  std::vector<FlowStringValue> Successors;
 
   std::vector<StringValue> Instructions;
 };
 
 template <> struct MappingTraits<MachineBasicBlock> {
   static void mapping(IO &YamlIO, MachineBasicBlock &MBB) {
     YamlIO.mapRequired("id", MBB.ID);
     YamlIO.mapOptional("name", MBB.Name,
                        std::string()); // Don't print out an empty name.
     YamlIO.mapOptional("alignment", MBB.Alignment);
     YamlIO.mapOptional("isLandingPad", MBB.IsLandingPad);
     YamlIO.mapOptional("addressTaken", MBB.AddressTaken);
+    YamlIO.mapOptional("successors", MBB.Successors);
     YamlIO.mapOptional("instructions", MBB.Instructions);
   }
 };
 
 } // end namespace yaml
 } // end namespace llvm
 
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::MachineBasicBlock)

lib/CodeGen/MIRParser/MIParser.h

 class SMDiagnostic;
 class SourceMgr;
 
 MachineInstr *
 parseMachineInstr(SourceMgr &SM, MachineFunction &MF, StringRef Src,
                   const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
                   const SlotMapping &IRSlots, SMDiagnostic &Error);
 
+MachineBasicBlock *
+parseMBBReference(SourceMgr &SM, MachineFunction &MF, StringRef Src,
+                  const DenseMap<unsigned, MachineBasicBlock *> &MBBMapping,
+                  const SlotMapping &IRSlots, SMDiagnostic &Error);
+
 } // end namespace llvm
 
 #endif

lib/CodeGen/MIRParser/MIParser.cpp

public:
   bool error(const Twine &Msg);
 
   /// Report an error at the given location with the given message.
   ///
   /// This function always return true.
   bool error(StringRef::iterator Loc, const Twine &Msg);
 
   MachineInstr *parse();
+  MachineBasicBlock *parseMBB();
 
   bool parseRegister(unsigned &Reg);
   bool parseRegisterOperand(MachineOperand &Dest, bool IsDef = false);
   bool parseImmediateOperand(MachineOperand &Dest);
+  MachineBasicBlock *parseMBBReference();
   bool parseMBBOperand(MachineOperand &Dest);
   bool parseGlobalAddressOperand(MachineOperand &Dest);
   bool parseMachineOperand(MachineOperand &Dest);
 
 private:
   /// Convert the integer literal in the current token into an unsigned integer.
   ///
   /// Return true if an error occurred.
MachineInstr *MIParser::parse() {
   // TODO: Determine the implicit behaviour when implicit register flags are
   // parsed.
   auto *MI = MF.CreateMachineInstr(MCID, DebugLoc(), /*NoImplicit=*/true);
   for (const auto &Operand : Operands)
     MI->addOperand(MF, Operand);
   return MI;
 }
 
+MachineBasicBlock *MIParser::parseMBB() {
+  lex();
+  if (Token.isNot(MIToken::MachineBasicBlock)) {
+    error("expected a machine basic block reference");
+    return nullptr;
+  }
+  auto *MBB = parseMBBReference();
+  lex();
+  if (Token.is(MIToken::Eof))
+    return MBB;
+  if (MBB) // Report this error only when MBB was parsed successfully.
+    error("expected end of string after the machine basic block reference");
+  return nullptr;
+}
+
 bool MIParser::parseInstruction(unsigned &OpCode) {
   if (Token.isNot(MIToken::Identifier))
     return error("expected a machine instruction");
   StringRef InstrName = Token.stringValue();
   if (parseInstrName(InstrName, OpCode))
     return error(Twine("unknown machine instruction name '") + InstrName + "'");
   lex();
   return false;
bool MIParser::getUnsigned(unsigned &Result) {
   const uint64_t Limit = uint64_t(std::numeric_limits<unsigned>::max()) + 1;
   uint64_t Val64 = Token.integerValue().getLimitedValue(Limit);
   if (Val64 == Limit)
     return error("expected 32-bit integer (too large)");
   Result = Val64;
   return false;
 }
 
-bool MIParser::parseMBBOperand(MachineOperand &Dest) {
+MachineBasicBlock *MIParser::parseMBBReference() {
   assert(Token.is(MIToken::MachineBasicBlock));
   unsigned Number;
   if (getUnsigned(Number))
-    return true;
+    return nullptr;
   auto MBBInfo = MBBSlots.find(Number);
-  if (MBBInfo == MBBSlots.end())
-    return error(Twine("use of undefined machine basic block #") +
-                 Twine(Number));
+  if (MBBInfo == MBBSlots.end()) {
+    error(Twine("use of undefined machine basic block #") + Twine(Number));
+    return nullptr;
+  }
   MachineBasicBlock *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() + "'");
+  if (!Token.stringValue().empty() && Token.stringValue() != MBB->getName()) {
+    error(Twine("the name of machine basic block #") + Twine(Number) +
+    return nullptr;
+  }
+  return MBB;
+}
+
+bool MIParser::parseMBBOperand(MachineOperand &Dest) {
+  if (auto *MBB = parseMBBReference()) {
     Dest = MachineOperand::CreateMBB(MBB);
     lex();
     return false;
   }
+  return true;
+}
 
 bool MIParser::parseGlobalAddressOperand(MachineOperand &Dest) {
   switch (Token.kind()) {
   case MIToken::NamedGlobalValue: {
     auto Name = Token.stringValue();
     const Module *M = MF.getFunction()->getParent();
     if (const auto *GV = M->getNamedValue(Name)) {
       Dest = MachineOperand::CreateGA(GV, /*Offset=*/0);

 }
 
 MachineInstr *
 llvm::parseMachineInstr(SourceMgr &SM, MachineFunction &MF, StringRef Src,
                         const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
                         const SlotMapping &IRSlots, SMDiagnostic &Error) {
   return MIParser(SM, MF, Error, Src, MBBSlots, IRSlots).parse();
 }
+
+MachineBasicBlock *llvm::parseMBBReference(
+    SourceMgr &SM, MachineFunction &MF, StringRef Src,
+    const DenseMap<unsigned, MachineBasicBlock *> &MBBMapping,
+    const SlotMapping &IRSlots, SMDiagnostic &Error) {
+  return MIParser(SM, MF, Error, Src, MBBMapping, IRSlots).parseMBB();
+}

lib/CodeGen/MIRParser/MIRParser.cpp

 bool MIRParserImpl::initializeMachineBasicBlock(
     MachineFunction &MF, MachineBasicBlock &MBB,
     const yaml::MachineBasicBlock &YamlMBB,
     const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots) {
   MBB.setAlignment(YamlMBB.Alignment);
   if (YamlMBB.AddressTaken)
     MBB.setHasAddressTaken();
   MBB.setIsLandingPad(YamlMBB.IsLandingPad);
+  SMDiagnostic Error;
+  // Parse the successors.
+  for (const auto &MBBSource : YamlMBB.Successors) {
+    if (auto *SuccMBB = parseMBBReference(SM, MF, MBBSource.Value, MBBSlots,
+                                          IRSlots, Error)) {
+      // TODO: Report an error when the successor is self.
+      // TODO: Report an error when adding the same successor more than once.
+      MBB.addSuccessor(SuccMBB);
+      continue;
+    }
+    reportDiagnostic(diagFromMIStringDiag(Error, MBBSource.SourceRange));
+    return true;
+  }
   // Parse the instructions.
   for (const auto &MISource : YamlMBB.Instructions) {
-    SMDiagnostic Error;
     if (auto *MI = parseMachineInstr(SM, MF, MISource.Value, MBBSlots, IRSlots,
                                      Error)) {
       MBB.insert(MBB.end(), MI);
       continue;
     }
     reportDiagnostic(diagFromMIStringDiag(Error, MISource.SourceRange));
     return true;
   }

lib/CodeGen/MIRPrinter.cpp

class MIPrinter {
   const DenseMap<const uint32_t *, unsigned> &RegisterMaskIds;
 
 public:
   MIPrinter(const Module &M, raw_ostream &OS,
             const DenseMap<const uint32_t *, unsigned> &RegisterMaskIds)
       : M(M), OS(OS), RegisterMaskIds(RegisterMaskIds) {}
 
   void print(const MachineInstr &MI);
+  void printMBBReference(const MachineBasicBlock &MBB);
   void print(const MachineOperand &Op, const TargetRegisterInfo *TRI);
 };
 
 } // end anonymous namespace
 
 namespace llvm {
 namespace yaml {
 
void MIRPrinter::convert(const Module &M, yaml::MachineBasicBlock &YamlMBB,
   // TODO: Serialize unnamed BB references.
   if (const auto *BB = MBB.getBasicBlock())
     YamlMBB.Name = BB->hasName() ? BB->getName() : "<unnamed bb>";
   else
     YamlMBB.Name = "";
   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());
   std::string Str;
   for (const auto &MI : MBB) {
     raw_string_ostream StrOS(Str);
     MIPrinter(M, StrOS, RegisterMaskIds).print(MI);
     YamlMBB.Instructions.push_back(StrOS.str());
static void printReg(unsigned Reg, raw_ostream &OS,
   if (!Reg)
     OS << '_';
   else if (Reg < TRI->getNumRegs())
     OS << '%' << StringRef(TRI->getName(Reg)).lower();
   else
     llvm_unreachable("Can't print this kind of register yet");
 }
 
+void MIPrinter::printMBBReference(const MachineBasicBlock &MBB) {
+  OS << "%bb." << MBB.getNumber();
+  if (const auto *BB = MBB.getBasicBlock()) {
+    if (BB->hasName())
+      OS << '.' << BB->getName();
+  }
+}
+
 void MIPrinter::print(const MachineOperand &Op, const TargetRegisterInfo *TRI) {
   switch (Op.getType()) {
   case MachineOperand::MO_Register:
     // TODO: Print register flags.
     printReg(Op.getReg(), OS, TRI);
     // TODO: Print sub register.
     break;
   case MachineOperand::MO_Immediate:
     OS << Op.getImm();
     break;
   case MachineOperand::MO_MachineBasicBlock:
-    OS << "%bb." << Op.getMBB()->getNumber();
-    if (const auto *BB = Op.getMBB()->getBasicBlock()) {
-      if (BB->hasName())
-        OS << '.' << BB->getName();
-    }
+    printMBBReference(*Op.getMBB());
     break;
   case MachineOperand::MO_GlobalAddress:
     // FIXME: Make this faster - print as operand will create a slot tracker to
     // print unnamed values for the whole module every time it's called, which
     // is inefficient.
     Op.getGlobal()->printAsOperand(OS, /*PrintType=*/false, &M);
     // TODO: Print offset and target flags.
     break;

test/CodeGen/MIR/expected-eof-after-successor-mbb.mir

+# RUN: not llc -start-after branch-folder -stop-after branch-folder -o /dev/null %s 2>&1 | FileCheck %s
+
+--- |
+
+  define i32 @foo(i32 %a) {
+  entry:
+    %0 = icmp sle i32 %a, 10
+    br i1 %0, label %less, label %exit
+
+  less:
+    ret i32 0
+
+  exit:
+    ret i32 %a
+  }
+
+...
+---
+name:            foo
+body:
+  - id:          0
+    name:        entry
+    # CHECK: [[@LINE+1]]:46: expected end of string after the machine basic block reference
+    successors:  [ '%bb.1.less', '%bb.2.exit 2' ]
+  - id:          1
+    name:        less
+  - id:          2
+    name:        exit
+...

test/CodeGen/MIR/expected-mbb-reference-for-successor-mbb.mir

+# RUN: not llc -start-after branch-folder -stop-after branch-folder -o /dev/null %s 2>&1 | FileCheck %s
+
+--- |
+
+  define i32 @foo(i32 %a) {
+  entry:
+    %0 = icmp sle i32 %a, 10
+    br i1 %0, label %less, label %exit
+
+  less:
+    ret i32 0
+
+  exit:
+    ret i32 %a
+  }
+
+...
+---
+name:            foo
+body:
+  - id:          0
+    name:        entry
+    # CHECK: [[@LINE+1]]:35: expected a machine basic block reference
+    successors:  [ '%bb.1.less', '2' ]
+  - id:          1
+    name:        less
+  - id:          2
+    name:        exit
+...

test/CodeGen/MIR/successor-basic-blocks.mir

+# RUN: llc -start-after branch-folder -stop-after branch-folder -o /dev/null %s | FileCheck %s
+# This test ensures that the MIR parser parses basic block successors correctly.
+
+--- |
+
+  define i32 @foo(i32 %a) {
+  entry:
+    %0 = icmp sle i32 %a, 10
+    br i1 %0, label %less, label %exit
+
+  less:
+    ret i32 0
+
+  exit:
+    ret i32 %a
+  }
+
+  define i32 @bar(i32 %a) {
+  entry:
+    %b = icmp sle i32 %a, 10
+    br i1 %b, label %0, label %1
+
+  ; <label>:0
+    ret i32 0
+
+  ; <label>:1
+    ret i32 %a
+  }
+
+...
+---
+name:            foo
+body:
+  # CHECK: name: entry
+  # CHECK: successors: [ '%bb.1.less', '%bb.2.exit' ]
+  # CHECK: name: less
+  - id:          0
+    name:        entry
+    successors:  [ '%bb.1.less', '%bb.2.exit' ]
+  - id:          1
+    name:        less
+  - id:          2
+    name:        exit
+...
+---
+name:            bar
+body:
+  # CHECK: name: bar
+  # CHECK: name: entry
+  # CHECK: successors: [ '%bb.1', '%bb.2' ]
+  # CHECK: id: 1
+  # CHECK: id: 2
+  - id:          0
+    name:        entry
+    successors:  [ '%bb.1', '%bb.2' ]
+  - id:          1
+  - id:          2
+...