[NFC] Refactor DXContainer to support more parts

[beanz]

Summary

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.

[lhames]

LGTM.

llvm/include/llvm/BinaryFormat/DXContainer.h

void swapBytes() {
     sys::swapByteOrder(ShaderKind);
     sys::swapByteOrder(Size);
     Bitcode.swapBytes();
   }
 };
 
 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
 
 #endif // LLVM_BINARYFORMAT_DXCONTAINER_H

llvm/include/llvm/BinaryFormat/DXContainerConstants.def

+
+#ifdef CONTAINER_PART
+CONTAINER_PART(DXIL)
+
+#undef CONTAINER_PART
+#endif 

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

struct DXILProgram {
   uint16_t DXILMajorVersion;
   uint16_t DXILMinorVersion;
   Optional<uint32_t> DXILOffset;
   Optional<uint32_t> DXILSize;
   Optional<std::vector<llvm::yaml::Hex8>> DXIL;
 };
 
 struct Part {
+  Part() = default;
+  Part(std::string N, uint32_t S) : Name(N), Size(S) {}
   std::string Name;
   uint32_t Size;
   Optional<DXILProgram> Program;
 };
 
 struct Object {
   FileHeader Header;
   std::vector<Part> Parts;

llvm/lib/BinaryFormat/CMakeLists.txt

 add_llvm_component_library(LLVMBinaryFormat
   AMDGPUMetadataVerifier.cpp
   COFF.cpp
   Dwarf.cpp
+  DXContainer.cpp
   ELF.cpp
   MachO.cpp
   Magic.cpp
   Minidump.cpp
   MsgPackDocument.cpp
   MsgPackDocumentYAML.cpp
   MsgPackReader.cpp
   MsgPackWriter.cpp

llvm/lib/BinaryFormat/DXContainer.cpp

+
+//===-- 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

for (uint32_t Part = 0; Part < Header.PartCount; ++Part) {
     // header. To prevent overflow, we subtract the part header size from the
     // buffer size, rather than adding to the offset. Since the file header is
     // larger than the part header we can't reach this code unless the buffer
     // is larger than the part header, so this can't underflow.
     if (PartOffset > Data.getBufferSize() - sizeof(dxbc::PartHeader))
       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;
+    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();
 }
 
 Expected<DXContainer> DXContainer::create(MemoryBufferRef Object) {
   DXContainer Container(Object);
   if (Error Err = Container.parseHeader())
     return std::move(Err);

llvm/tools/obj2yaml/dxcontainer2yaml.cpp

dumpDXContainer(MemoryBufferRef Source) {
   Obj->Header.Version.Major = Container.getHeader().Version.Major;
   Obj->Header.Version.Minor = Container.getHeader().Version.Minor;
   Obj->Header.FileSize = Container.getHeader().FileSize;
   Obj->Header.PartCount = Container.getHeader().PartCount;
 
   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.MinorVersion,
+          DXIL->first.Bitcode.Offset,
+              DXIL->second, DXIL->second + DXIL->first.Bitcode.Size)};
+      break;
+    }
+    case dxbc::PartType::Unknown:
+      break;
     }
   }
 
   return Obj.release();
 }
 
 llvm::Error dxcontainer2yaml(llvm::raw_ostream &Out,
                              llvm::MemoryBufferRef Source) {