[NFC] Refactoring DXContainerYaml Root Parameter representation

llvm/include/llvm/ObjectYAML/DXContainerYAML.h

@@ -90,29 +90,16 @@ struct RootDescriptorYaml {
 #include "llvm/BinaryFormat/DXContainerConstants.def"
 };
 
+using ParameterData = std::variant<RootConstantsYaml, RootDescriptorYaml>;
+
 struct RootParameterYamlDesc {
   uint32_t Type;
   uint32_t Visibility;
   uint32_t Offset;
-  RootParameterYamlDesc() {};
-  RootParameterYamlDesc(uint32_t T) : Type(T) {
-    switch (T) {
-
-    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
-      Constants = RootConstantsYaml();
-      break;
-    case llvm::to_underlying(dxbc::RootParameterType::CBV):
-    case llvm::to_underlying(dxbc::RootParameterType::SRV):
-    case llvm::to_underlying(dxbc::RootParameterType::UAV):
-      Descriptor = RootDescriptorYaml();
-      break;
-    }
-  }
+  ParameterData Data;
 
-  union {
-    RootConstantsYaml Constants;
-    RootDescriptorYaml Descriptor;
-  };
+  RootParameterYamlDesc() {};
+  RootParameterYamlDesc(uint32_t T) : Type(T) {}
 };
 
 struct RootSignatureYamlDesc {

llvm/lib/ObjectYAML/DXContainerEmitter.cpp

@@ -274,28 +274,26 @@ void DXContainerWriter::writeParts(raw_ostream &OS) {
       RS.StaticSamplersOffset = P.RootSignature->StaticSamplersOffset;
 
       for (const auto &Param : P.RootSignature->Parameters) {
-        dxbc::RootParameterHeader Header{Param.Type, Param.Visibility,
-                                         Param.Offset};
+        auto Header = dxbc::RootParameterHeader{Param.Type, Param.Visibility,
+                                                Param.Offset};
 
-        switch (Param.Type) {
-        case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
+        if (auto *ConstantYaml =
+                std::get_if<DXContainerYAML::RootConstantsYaml>(&Param.Data)) {
           dxbc::RootConstants Constants;
-          Constants.Num32BitValues = Param.Constants.Num32BitValues;
-          Constants.RegisterSpace = Param.Constants.RegisterSpace;
-          Constants.ShaderRegister = Param.Constants.ShaderRegister;
+          Constants.Num32BitValues = ConstantYaml->Num32BitValues;
+          Constants.RegisterSpace = ConstantYaml->RegisterSpace;
+          Constants.ShaderRegister = ConstantYaml->ShaderRegister;
           RS.ParametersContainer.addParameter(Header, Constants);
-          break;
-        case llvm::to_underlying(dxbc::RootParameterType::SRV):
-        case llvm::to_underlying(dxbc::RootParameterType::UAV):
-        case llvm::to_underlying(dxbc::RootParameterType::CBV):
+        } else if (auto *DescriptorYaml =
+                       std::get_if<DXContainerYAML::RootDescriptorYaml>(
+                           &Param.Data)) {
           dxbc::RTS0::v2::RootDescriptor Descriptor;
-          Descriptor.RegisterSpace = Param.Descriptor.RegisterSpace;
-          Descriptor.ShaderRegister = Param.Descriptor.ShaderRegister;
+          Descriptor.RegisterSpace = DescriptorYaml->RegisterSpace;
+          Descriptor.ShaderRegister = DescriptorYaml->ShaderRegister;
           if (RS.Version > 1)
-            Descriptor.Flags = Param.Descriptor.getEncodedFlags();
+            Descriptor.Flags = DescriptorYaml->getEncodedFlags();
           RS.ParametersContainer.addParameter(Header, Descriptor);
-          break;
-        default:
+        } else {
           // Handling invalid parameter type edge case. We intentionally let
           // obj2yaml/yaml2obj parse and emit invalid dxcontainer data, in order
           // for that to be used as a testing tool more effectively.

llvm/lib/ObjectYAML/DXContainerYAML.cpp

@@ -75,26 +75,29 @@ DXContainerYAML::RootSignatureYamlDesc::create(
         return std::move(E);
 
       auto Constants = *ConstantsOrErr;
-
-      NewP.Constants.Num32BitValues = Constants.Num32BitValues;
-      NewP.Constants.ShaderRegister = Constants.ShaderRegister;
-      NewP.Constants.RegisterSpace = Constants.RegisterSpace;
+      RootConstantsYaml ConstantYaml;
+      ConstantYaml.Num32BitValues = Constants.Num32BitValues;
+      ConstantYaml.ShaderRegister = Constants.ShaderRegister;
+      ConstantYaml.RegisterSpace = Constants.RegisterSpace;
+      NewP.Data = ConstantYaml;
     } else if (auto *RDV =
                    dyn_cast<object::DirectX::RootDescriptorView>(&ParamView)) {
       llvm::Expected<dxbc::RTS0::v2::RootDescriptor> DescriptorOrErr =
           RDV->read(Version);
       if (Error E = DescriptorOrErr.takeError())
         return std::move(E);
       auto Descriptor = *DescriptorOrErr;
-      NewP.Descriptor.ShaderRegister = Descriptor.ShaderRegister;
-      NewP.Descriptor.RegisterSpace = Descriptor.RegisterSpace;
+      RootDescriptorYaml YamlDescriptor;
+      YamlDescriptor.ShaderRegister = Descriptor.ShaderRegister;
+      YamlDescriptor.RegisterSpace = Descriptor.RegisterSpace;
       if (Version > 1) {
 #define ROOT_DESCRIPTOR_FLAG(Num, Val)                                         \
-  NewP.Descriptor.Val =                                                        \
+  YamlDescriptor.Val =                                                         \
       (Descriptor.Flags &                                                      \
        llvm::to_underlying(dxbc::RootDescriptorFlag::Val)) > 0;
 #include "llvm/BinaryFormat/DXContainerConstants.def"
       }
+      NewP.Data = YamlDescriptor;
     }
 
     RootSigDesc.Parameters.push_back(NewP);
@@ -316,14 +319,24 @@ void MappingTraits<llvm::DXContainerYAML::RootParameterYamlDesc>::mapping(
   IO.mapRequired("ShaderVisibility", P.Visibility);
 
   switch (P.Type) {
-  case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
-    IO.mapRequired("Constants", P.Constants);
-    break;
+  case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+    DXContainerYAML::RootConstantsYaml Constants;
+    if (IO.outputting())
+      Constants = std::get<DXContainerYAML::RootConstantsYaml>(P.Data);
+    IO.mapRequired("Constants", Constants);
+    if (!IO.outputting())
+      P.Data = Constants;
+  } break;
   case llvm::to_underlying(dxbc::RootParameterType::CBV):
   case llvm::to_underlying(dxbc::RootParameterType::SRV):
-  case llvm::to_underlying(dxbc::RootParameterType::UAV):
-    IO.mapRequired("Descriptor", P.Descriptor);
-    break;
+  case llvm::to_underlying(dxbc::RootParameterType::UAV): {
+    DXContainerYAML::RootDescriptorYaml Descriptor;
+    if (IO.outputting())
+      Descriptor = std::get<DXContainerYAML::RootDescriptorYaml>(P.Data);
+    IO.mapRequired("Descriptor", Descriptor);
+    if (!IO.outputting())
+      P.Data = Descriptor;
+  } break;
   }
 }