[llvm] [NFC] Refactoring MCDXBC to support out of order storage of root parameters (PR #137284)

[Justin Bogner via llvm-commits]

================
@@ -9,18 +9,81 @@
 #include "llvm/BinaryFormat/DXContainer.h"
 #include <cstdint>
 #include <limits>
+#include <variant>
 
 namespace llvm {
 
 class raw_ostream;
 namespace mcdxbc {
 
-struct RootParameter {
+struct RootParameterInfo {
   dxbc::RootParameterHeader Header;
-  union {
-    dxbc::RootConstants Constants;
-    dxbc::RTS0::v2::RootDescriptor Descriptor;
-  };
+  size_t Location;
+
+  RootParameterInfo() = default;
+
+  RootParameterInfo(dxbc::RootParameterHeader H, size_t L)
+      : Header(H), Location(L) {}
+};
+
+using RootDescriptor = std::variant<dxbc::RTS0::v1::RootDescriptor,
+                                    dxbc::RTS0::v2::RootDescriptor>;
+using ParametersView = std::variant<const dxbc::RootConstants *,
+                                    const dxbc::RTS0::v1::RootDescriptor *,
+                                    const dxbc::RTS0::v2::RootDescriptor *>;
+struct RootParametersContainer {
+  SmallVector<RootParameterInfo> ParametersInfo;
+
+  SmallVector<dxbc::RootConstants> Constants;
+  SmallVector<RootDescriptor> Descriptors;
+
+  void addInfo(dxbc::RootParameterHeader H, size_t L) {
+    ParametersInfo.push_back(RootParameterInfo(H, L));
+  }
+
+  void addParameter(dxbc::RootParameterHeader H, dxbc::RootConstants C) {
+    addInfo(H, Constants.size());
+    Constants.push_back(C);
+  }
+
+  void addParameter(dxbc::RootParameterHeader H,
+                    dxbc::RTS0::v1::RootDescriptor D) {
+    addInfo(H, Descriptors.size());
+    Descriptors.push_back(D);
+  }
+
+  void addParameter(dxbc::RootParameterHeader H,
+                    dxbc::RTS0::v2::RootDescriptor D) {
+    addInfo(H, Descriptors.size());
+    Descriptors.push_back(D);
+  }
+
+  std::optional<ParametersView> getParameter(const RootParameterInfo *H) const {
+    switch (H->Header.ParameterType) {
+    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit):
+      return &Constants[H->Location];
+    case llvm::to_underlying(dxbc::RootParameterType::CBV):
+    case llvm::to_underlying(dxbc::RootParameterType::SRV):
+    case llvm::to_underlying(dxbc::RootParameterType::UAV):
+      const RootDescriptor &VersionedParam = Descriptors[H->Location];
+      if (std::holds_alternative<dxbc::RTS0::v1::RootDescriptor>(
+              VersionedParam)) {
+        return &std::get<dxbc::RTS0::v1::RootDescriptor>(VersionedParam);
+      }
+      return &std::get<dxbc::RTS0::v2::RootDescriptor>(VersionedParam);
+    }
+
+    return std::nullopt;
+  }
----------------
bogner wrote:

I'm not really convinced that the abstraction that the `std::variant` and the generic `getParameter` are providing are really helping much here. Consider this simplified API:
```c++
  const std::pair<uint32_t, uint32_t>
  getTypeAndLocForParameter(uint32_t Index) const {
    const RootParameterInfo &Info = ParametersInfo[Index];
    return {Info.Header.ParameterType, Info.Location};
  }

  const dxbc::RootConstants &getConstant(size_t Index) const {
    return Constants[Index];
  }

  const dxbc::RTS0::v2::RootDescriptor &getRootDescriptor(size_t Index) const {
    return Descriptors[Index];
  }
```

The logic to use this is more or less the same - instead of the `holds_alternative` type checks we simply check against the type enum we already have:

```diff
-    auto P = ParametersContainer.getParameter(ParametersContainer[I]);
-    if (std::holds_alternative<const dxbc::RootConstants *>(P.value())) {
-      auto *Constants = std::get<const dxbc::RootConstants *>(P.value());
-      support::endian::write(BOS, Constants->ShaderRegister,
-                             llvm::endianness::little);
+    const auto &[Type, Loc] = ParametersContainer.getTypeAndLocForParameter(I);
+    switch (Type) {
+    case llvm::to_underlying(dxbc::RootParameterType::Constants32Bit): {
+      const dxbc::RootConstants &Constants =
+          ParametersContainer.getConstant(Loc);
+      support::endian::write(BOS, Constants.ShaderRegister,
+                             llvm::endianness::little);
```

I think `std::variant` has its uses when we need something akin to a type safe union or we want to use visitor patterns to handle a large number of cases, but if we're just going to switch over the types anyway I think it just adds a layer of abstraction that needs to be looked through when reading the code to understand it.

https://github.com/llvm/llvm-project/pull/137284