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;
 
 bool parseMachineInstr(MachineInstr *&MI, SourceMgr &SM, MachineFunction &MF,
                        StringRef Src,
                        const DenseMap<unsigned, MachineBasicBlock *> &MBBSlots,
                        const SlotMapping &IRSlots, SMDiagnostic &Error);
 
+bool parseMBBReference(
+    MachineBasicBlock *&MBB, 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);
 
   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);
 
 private:
   /// Convert the integer literal in the current token into an unsigned integer.
   ///
   /// Return true if an error occurred.
bool MIParser::parse(MachineInstr *&MI) {
   // TODO: Determine the implicit behaviour when implicit register flags are
   // parsed.
   MI = MF.CreateMachineInstr(MCID, DebugLoc(), /*NoImplicit=*/true);
   for (const auto &Operand : Operands)
     MI->addOperand(MF, Operand);
   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");
   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) {
+bool MIParser::parseMBBReference(MachineBasicBlock *&MBB) {
   assert(Token.is(MIToken::MachineBasicBlock));
   unsigned Number;
   if (getUnsigned(Number))
     return true;
   auto MBBInfo = MBBSlots.find(Number);
   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;
 }
 
 bool MIParser::parseGlobalAddressOperand(MachineOperand &Dest) {
   switch (Token.kind()) {
   case MIToken::NamedGlobalValue: {

 }
 
 bool llvm::parseMachineInstr(
     MachineInstr *&MI, 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(MI);
 }
+
+bool llvm::parseMBBReference(
+    MachineBasicBlock *&MBB, 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(MBB);
+}

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) {
+    MachineBasicBlock *SuccMBB = nullptr;
+    if (parseMBBReference(SuccMBB, SM, MF, MBBSource.Value, MBBSlots, IRSlots,
+                          Error))
+      return error(Error, MBBSource.SourceRange);
+    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);
     MBB.insert(MBB.end(), MI);
   }
   return false;
 }
 

lib/CodeGen/MIRPrinter.cpp

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());

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
+...