[NFC] Refactor DXContainer to support more parts

This patch refactors some of the DXContainer Object and YAML code to
make it easier to add more part parsing.

DXContainer has a whole bunch of constant values, so I've added a
DXContainerConstants.def file which will grow with constant
definitions, but starts with just part identifiers. I've also added a
utility to parse the part magic string into an enum, and converted the
code to use that utility and the enum instead of the part literal
string.

Reviewed By: lhames

Differential Revision: https://reviews.llvm.org/D133980

llvm/include/llvm/BinaryFormat/DXContainer.h

@@ -125,6 +125,14 @@ struct ProgramHeader {
 
 static_assert(sizeof(ProgramHeader) == 24, "ProgramHeader Size incorrect!");
 
+#define CONTAINER_PART(Part) Part,
+enum class PartType {
+  Unknown = 0,
+#include "DXContainerConstants.def"
+};
+
+PartType parsePartType(StringRef S);
+
 } // namespace dxbc
 } // namespace llvm
 

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

@@ -0,0 +1,6 @@
+
+#ifdef CONTAINER_PART
+CONTAINER_PART(DXIL)
+
+#undef CONTAINER_PART
+#endif 

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -54,6 +54,8 @@ struct DXILProgram {
 };
 
 struct Part {
+  Part() = default;
+  Part(std::string N, uint32_t S) : Name(N), Size(S) {}
   std::string Name;
   uint32_t Size;
   Optional<DXILProgram> Program;

llvm/lib/BinaryFormat/CMakeLists.txt

@@ -2,6 +2,7 @@ add_llvm_component_library(LLVMBinaryFormat
   AMDGPUMetadataVerifier.cpp
   COFF.cpp
   Dwarf.cpp
+  DXContainer.cpp
   ELF.cpp
   MachO.cpp
   Magic.cpp

llvm/lib/BinaryFormat/DXContainer.cpp

@@ -0,0 +1,25 @@
+
+//===-- llvm/BinaryFormat/DXContainer.cpp - DXContainer Utils ----*- C++-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains utility functions for working with DXContainers.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/ADT/StringSwitch.h"
+
+using namespace llvm;
+using namespace llvm;
+
+dxbc::PartType dxbc::parsePartType(StringRef S) {
+#define CONTAINER_PART(PartName) .Case(#PartName, PartType::PartName)
+  return StringSwitch<dxbc::PartType>(S)
+#include "llvm/BinaryFormat/DXContainerConstants.def"
+      .Default(dxbc::PartType::Unknown);
+}

llvm/lib/Object/DXContainer.cpp

@@ -81,11 +81,16 @@ Error DXContainer::parsePartOffsets() {
       return parseFailed("Part offset points beyond boundary of the file");
     PartOffsets.push_back(PartOffset);
 
-    // If this isn't a dxil part stop here...
-    if (Data.getBuffer().substr(PartOffset, 4) != "DXIL")
-      continue;
-    if (Error Err = parseDXILHeader(PartOffset + sizeof(dxbc::PartHeader)))
-      return Err;
+    dxbc::PartType PT =
+        dxbc::parsePartType(Data.getBuffer().substr(PartOffset, 4));
+    switch (PT) {
+    case dxbc::PartType::DXIL:
+      if (Error Err = parseDXILHeader(PartOffset + sizeof(dxbc::PartHeader)))
+        return Err;
+      break;
+    case dxbc::PartType::Unknown:
+      break;
+    }
   }
   return Error::success();
 }

llvm/tools/obj2yaml/dxcontainer2yaml.cpp

@@ -38,23 +38,30 @@ dumpDXContainer(MemoryBufferRef Source) {
   Obj->Header.PartOffsets = std::vector<uint32_t>();
   for (const auto P : Container) {
     Obj->Header.PartOffsets->push_back(P.Offset);
-    if (P.Part.getName() == "DXIL") {
+    Obj->Parts.push_back(
+        DXContainerYAML::Part(P.Part.getName().str(), P.Part.Size));
+    DXContainerYAML::Part &NewPart = Obj->Parts.back();
+    dxbc::PartType PT = dxbc::parsePartType(P.Part.getName());
+    switch (PT) {
+    case dxbc::PartType::DXIL: {
       Optional<DXContainer::DXILData> DXIL = Container.getDXIL();
       assert(DXIL && "Since we are iterating and found a DXIL part, "
                      "this should never not have a value");
-      Obj->Parts.push_back(DXContainerYAML::Part{
-          P.Part.getName().str(), P.Part.Size,
-          DXContainerYAML::DXILProgram{
-              DXIL->first.MajorVersion, DXIL->first.MinorVersion,
-              DXIL->first.ShaderKind, DXIL->first.Size,
-              DXIL->first.Bitcode.MajorVersion,
-              DXIL->first.Bitcode.MinorVersion, DXIL->first.Bitcode.Offset,
-              DXIL->first.Bitcode.Size,
-              std::vector<llvm::yaml::Hex8>(
-                  DXIL->second, DXIL->second + DXIL->first.Bitcode.Size)}});
-    } else {
-      Obj->Parts.push_back(
-          DXContainerYAML::Part{P.Part.getName().str(), P.Part.Size, None});
+      NewPart.Program = DXContainerYAML::DXILProgram{
+          DXIL->first.MajorVersion,
+          DXIL->first.MinorVersion,
+          DXIL->first.ShaderKind,
+          DXIL->first.Size,
+          DXIL->first.Bitcode.MajorVersion,
+          DXIL->first.Bitcode.MinorVersion,
+          DXIL->first.Bitcode.Offset,
+          DXIL->first.Bitcode.Size,
+          std::vector<llvm::yaml::Hex8>(
+              DXIL->second, DXIL->second + DXIL->first.Bitcode.Size)};
+      break;
+    }
+    case dxbc::PartType::Unknown:
+      break;
     }
   }