MIR Serialization: Serialize the jump table info.

The jump table info is serialized using a YAML mapping that contains its kind
and a YAML sequence of jump table entries. A jump table entry is a YAML mapping
that has an ID and an inline YAML sequence of machine basic block references.

The testcase 'CodeGen/MIR/X86/jump-table-info.mir' doesn't have any instructions
because one of them contains a jump table index operand. The jump table index
operands will be serialized in a follow up patch, and the appropriate
instructions will be added to this testcase.

Reviewers: Duncan P. N. Exon Smith
llvm-svn: 242357

include/llvm/CodeGen/MIRYamlMapping.h

@@ -19,6 +19,7 @@
 #define LLVM_LIB_CODEGEN_MIRYAMLMAPPING_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/Support/YAMLTraits.h"
 #include <vector>
 
@@ -72,6 +73,22 @@ template <> struct ScalarTraits<FlowStringValue> {
   static bool mustQuote(StringRef Scalar) { return needsQuotes(Scalar); }
 };
 
+template <> struct ScalarEnumerationTraits<MachineJumpTableInfo::JTEntryKind> {
+  static void enumeration(yaml::IO &IO,
+                          MachineJumpTableInfo::JTEntryKind &EntryKind) {
+    IO.enumCase(EntryKind, "block-address",
+                MachineJumpTableInfo::EK_BlockAddress);
+    IO.enumCase(EntryKind, "gp-rel64-block-address",
+                MachineJumpTableInfo::EK_GPRel64BlockAddress);
+    IO.enumCase(EntryKind, "gp-rel32-block-address",
+                MachineJumpTableInfo::EK_GPRel32BlockAddress);
+    IO.enumCase(EntryKind, "label-difference32",
+                MachineJumpTableInfo::EK_LabelDifference32);
+    IO.enumCase(EntryKind, "inline", MachineJumpTableInfo::EK_Inline);
+    IO.enumCase(EntryKind, "custom32", MachineJumpTableInfo::EK_Custom32);
+  }
+};
+
 } // end namespace yaml
 } // end namespace llvm
 
@@ -206,6 +223,23 @@ template <> struct MappingTraits<FixedMachineStackObject> {
   static const bool flow = true;
 };
 
+struct MachineJumpTable {
+  struct Entry {
+    unsigned ID;
+    std::vector<FlowStringValue> Blocks;
+  };
+
+  MachineJumpTableInfo::JTEntryKind Kind = MachineJumpTableInfo::EK_Custom32;
+  std::vector<Entry> Entries;
+};
+
+template <> struct MappingTraits<MachineJumpTable::Entry> {
+  static void mapping(IO &YamlIO, MachineJumpTable::Entry &Entry) {
+    YamlIO.mapRequired("id", Entry.ID);
+    YamlIO.mapOptional("blocks", Entry.Blocks);
+  }
+};
+
 } // end namespace yaml
 } // end namespace llvm
 
@@ -213,10 +247,18 @@ LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::VirtualRegisterDefinition)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::MachineBasicBlock)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::MachineStackObject)
 LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::FixedMachineStackObject)
+LLVM_YAML_IS_SEQUENCE_VECTOR(llvm::yaml::MachineJumpTable::Entry)
 
 namespace llvm {
 namespace yaml {
 
+template <> struct MappingTraits<MachineJumpTable> {
+  static void mapping(IO &YamlIO, MachineJumpTable &JT) {
+    YamlIO.mapRequired("kind", JT.Kind);
+    YamlIO.mapOptional("entries", JT.Entries);
+  }
+};
+
 /// Serializable representation of MachineFrameInfo.
 ///
 /// Doesn't serialize attributes like 'StackAlignment', 'IsStackRealignable' and
@@ -278,6 +320,7 @@ struct MachineFunction {
   MachineFrameInfo FrameInfo;
   std::vector<FixedMachineStackObject> FixedStackObjects;
   std::vector<MachineStackObject> StackObjects;
+  MachineJumpTable JumpTableInfo;
 
   std::vector<MachineBasicBlock> BasicBlocks;
 };
@@ -295,6 +338,8 @@ template <> struct MappingTraits<MachineFunction> {
     YamlIO.mapOptional("frameInfo", MF.FrameInfo);
     YamlIO.mapOptional("fixedStack", MF.FixedStackObjects);
     YamlIO.mapOptional("stack", MF.StackObjects);
+    if (!YamlIO.outputting() || !MF.JumpTableInfo.Entries.empty())
+      YamlIO.mapOptional("jumpTable", MF.JumpTableInfo);
     YamlIO.mapOptional("body", MF.BasicBlocks);
   }
 };

lib/CodeGen/MIRParser/MIRParser.cpp

@@ -111,6 +111,10 @@ public:
   bool initializeFrameInfo(const Function &F, MachineFrameInfo &MFI,
                            const yaml::MachineFunction &YamlMF);
 
+  bool initializeJumpTableInfo(MachineFunction &MF,
+                               const yaml::MachineJumpTable &YamlJTI,
+                               const PerFunctionMIParsingState &PFS);
+
 private:
   /// Return a MIR diagnostic converted from an MI string diagnostic.
   SMDiagnostic diagFromMIStringDiag(const SMDiagnostic &Error,
@@ -292,6 +296,11 @@ bool MIRParserImpl::initializeMachineFunction(MachineFunction &MF) {
   if (YamlMF.BasicBlocks.empty())
     return error(Twine("machine function '") + Twine(MF.getName()) +
                  "' requires at least one machine basic block in its body");
+  // Initialize the jump table after creating all the MBBs so that the MBB
+  // references can be resolved.
+  if (!YamlMF.JumpTableInfo.Entries.empty() &&
+      initializeJumpTableInfo(MF, YamlMF.JumpTableInfo, PFS))
+    return true;
   // Initialize the machine basic blocks after creating them all so that the
   // machine instructions parser can resolve the MBB references.
   unsigned I = 0;
@@ -425,6 +434,24 @@ bool MIRParserImpl::initializeFrameInfo(const Function &F,
   return false;
 }
 
+bool MIRParserImpl::initializeJumpTableInfo(
+    MachineFunction &MF, const yaml::MachineJumpTable &YamlJTI,
+    const PerFunctionMIParsingState &PFS) {
+  MachineJumpTableInfo *JTI = MF.getOrCreateJumpTableInfo(YamlJTI.Kind);
+  SMDiagnostic Error;
+  for (const auto &Entry : YamlJTI.Entries) {
+    std::vector<MachineBasicBlock *> Blocks;
+    for (const auto &MBBSource : Entry.Blocks) {
+      MachineBasicBlock *MBB = nullptr;
+      if (parseMBBReference(MBB, SM, MF, MBBSource.Value, PFS, IRSlots, Error))
+        return error(Error, MBBSource.SourceRange);
+      Blocks.push_back(MBB);
+    }
+    JTI->createJumpTableIndex(Blocks);
+  }
+  return false;
+}
+
 SMDiagnostic MIRParserImpl::diagFromMIStringDiag(const SMDiagnostic &Error,
                                                  SMRange SourceRange) {
   assert(SourceRange.isValid() && "Invalid source range");

lib/CodeGen/MIRPrinter.cpp

@@ -46,6 +46,8 @@ public:
   void convert(yaml::MachineFunction &MF, const MachineRegisterInfo &RegInfo,
                const TargetRegisterInfo *TRI);
   void convert(yaml::MachineFrameInfo &YamlMFI, const MachineFrameInfo &MFI);
+  void convert(ModuleSlotTracker &MST, yaml::MachineJumpTable &YamlJTI,
+               const MachineJumpTableInfo &JTI);
   void convert(ModuleSlotTracker &MST, yaml::MachineBasicBlock &YamlMBB,
                const MachineBasicBlock &MBB);
   void convertStackObjects(yaml::MachineFunction &MF,
@@ -116,8 +118,10 @@ void MIRPrinter::print(const MachineFunction &MF) {
   convert(YamlMF.FrameInfo, *MF.getFrameInfo());
   convertStackObjects(YamlMF, *MF.getFrameInfo());
 
-  int I = 0;
   ModuleSlotTracker MST(MF.getFunction()->getParent());
+  if (const auto *JumpTableInfo = MF.getJumpTableInfo())
+    convert(MST, YamlMF.JumpTableInfo, *JumpTableInfo);
+  int I = 0;
   for (const auto &MBB : MF) {
     // TODO: Allow printing of non sequentially numbered MBBs.
     // This is currently needed as the basic block references get their index
@@ -218,6 +222,25 @@ void MIRPrinter::convertStackObjects(yaml::MachineFunction &MF,
   }
 }
 
+void MIRPrinter::convert(ModuleSlotTracker &MST,
+                         yaml::MachineJumpTable &YamlJTI,
+                         const MachineJumpTableInfo &JTI) {
+  YamlJTI.Kind = JTI.getEntryKind();
+  unsigned ID = 0;
+  for (const auto &Table : JTI.getJumpTables()) {
+    std::string Str;
+    yaml::MachineJumpTable::Entry Entry;
+    Entry.ID = ID++;
+    for (const auto *MBB : Table.MBBs) {
+      raw_string_ostream StrOS(Str);
+      MIPrinter(StrOS, MST, RegisterMaskIds).printMBBReference(*MBB);
+      Entry.Blocks.push_back(StrOS.str());
+      Str.clear();
+    }
+    YamlJTI.Entries.push_back(Entry);
+  }
+}
+
 void MIRPrinter::convert(ModuleSlotTracker &MST,
                          yaml::MachineBasicBlock &YamlMBB,
                          const MachineBasicBlock &MBB) {

test/CodeGen/MIR/X86/jump-table-info.mir

@@ -0,0 +1,62 @@
+# RUN: llc -start-after branch-folder -stop-after branch-folder -o /dev/null %s | FileCheck %s
+# This test ensures that the MIR parser parses the jump table info correctly.
+
+--- |
+
+  define i32 @test_jumptable(i32 %in) {
+  entry:
+    switch i32 %in, label %def [
+      i32 0, label %lbl1
+      i32 1, label %lbl2
+      i32 2, label %lbl3
+      i32 3, label %lbl4
+    ]
+
+  def:
+    ret i32 0
+
+  lbl1:
+    ret i32 1
+
+  lbl2:
+    ret i32 2
+
+  lbl3:
+    ret i32 4
+
+  lbl4:
+    ret i32 8
+  }
+
+...
+---
+name:            test_jumptable
+# CHECK:      jumpTable:
+# CHECK-NEXT: kind: label-difference32
+# CHECK-NEXT: entries:
+# CHECK-NEXT: - id: 0
+# CHECK-NEXT: blocks: [ '%bb.3.lbl1', '%bb.4.lbl2', '%bb.5.lbl3', '%bb.6.lbl4' ]
+# CHECK_NEXT: body:
+jumpTable:
+  kind:          label-difference32
+  entries:
+    - id:        0
+      blocks:    [ '%bb.3.lbl1', '%bb.4.lbl2', '%bb.5.lbl3', '%bb.6.lbl4' ]
+body:
+  - id:          0
+    name:        entry
+    successors:  [ '%bb.2.def', '%bb.1.entry' ]
+  - id:          1
+    name:        entry
+    successors:  [ '%bb.3.lbl1', '%bb.4.lbl2', '%bb.5.lbl3', '%bb.6.lbl4' ]
+  - id:          2
+    name:        def
+  - id:          3
+    name:        lbl1
+  - id:          4
+    name:        lbl2
+  - id:          5
+    name:        lbl3
+  - id:          6
+    name:        lbl4
+...

test/CodeGen/MIR/invalid-jump-table-kind.mir

@@ -0,0 +1,54 @@
+# RUN: not llc -start-after branch-folder -stop-after branch-folder -o /dev/null %s 2>&1 | FileCheck %s
+
+--- |
+
+  define i32 @test_jumptable(i32 %in) {
+  entry:
+    switch i32 %in, label %def [
+      i32 0, label %lbl1
+      i32 1, label %lbl2
+      i32 2, label %lbl3
+      i32 3, label %lbl4
+    ]
+
+  def:
+    ret i32 0
+
+  lbl1:
+    ret i32 1
+
+  lbl2:
+    ret i32 2
+
+  lbl3:
+    ret i32 4
+
+  lbl4:
+    ret i32 8
+  }
+
+...
+---
+name:            test_jumptable
+jumpTable:
+  # CHECK: [[@LINE+1]]:18: unknown enumerated scalar
+  kind:          switch
+  entries:
+    - id:        0
+      blocks:    [ '%bb.3.lbl1', '%bb.4.lbl2', '%bb.5.lbl3', '%bb.6.lbl4' ]
+body:
+  - id:          0
+    name:        entry
+  - id:          1
+    name:        entry
+  - id:          2
+    name:        def
+  - id:          3
+    name:        lbl1
+  - id:          4
+    name:        lbl2
+  - id:          5
+    name:        lbl3
+  - id:          6
+    name:        lbl4
+...