[FIRRTL] Add some more aggregate type manipulation utils

include/circt/Dialect/FIRRTL/FIRRTLTypes.h

@@ -109,9 +109,14 @@ public:
   /// nested under another type.
   unsigned getMaxFieldID();
 
-  /// Returns the effective field id when treating the index field as the root
-  /// of the type.  Essentially maps a fieldID to a fieldID after a subfield op.
-  /// Returns the new id and whether the id is in the given child.
+  /// Get the sub-type of a type for a field ID.  This is the identity function
+  /// for ground types.
+  FIRRTLType getSubTypeByFieldID(unsigned fieldID);
+
+  /// Returns the effective field id when treating the index field as the
+  /// root of the type.  Essentially maps a fieldID to a fieldID after a
+  /// subfield op. Returns the new id and whether the id is in the given
+  /// child.
   std::pair<unsigned, bool> rootChildFieldID(unsigned fieldID, unsigned index);
 
 protected:
@@ -305,6 +310,9 @@ public:
   /// Look up an element type by name.
   FIRRTLType getElementType(StringRef name);
 
+  /// Look up an element type by index.
+  FIRRTLType getElementType(size_t index);
+
   /// Return the recursive properties of the type.
   RecursiveTypeProperties getRecursiveTypeProperties();
 
@@ -323,6 +331,8 @@ public:
   /// index of the parent field.
   unsigned getIndexForFieldID(unsigned fieldID);
 
+  FIRRTLType getSubTypeByFieldID(unsigned fieldID);
+
   /// Get the maximum field ID in this bundle.  This is helpful for constructing
   /// field IDs when this BundleType is nested in another aggregate type.
   unsigned getMaxFieldID();
@@ -364,6 +374,8 @@ public:
   /// the index of the parent element.
   unsigned getIndexForFieldID(unsigned fieldID);
 
+  FIRRTLType getSubTypeByFieldID(unsigned fieldID);
+
   /// Get the maximum field ID in this vector.  This is helpful for constructing
   /// field IDs when this VectorType is nested in another aggregate type.
   unsigned getMaxFieldID();

lib/Dialect/FIRRTL/FIRRTLTypes.cpp

@@ -398,6 +398,18 @@ unsigned FIRRTLType::getMaxFieldID() {
       });
 }
 
+FIRRTLType FIRRTLType::getSubTypeByFieldID(unsigned fieldID) {
+  return TypeSwitch<FIRRTLType, FIRRTLType>(*this)
+      .Case<AnalogType, ClockType, ResetType, AsyncResetType, SIntType,
+            UIntType>([](FIRRTLType t) { return t; })
+      .Case<BundleType, FVectorType>(
+          [fieldID](auto type) { return type.getSubTypeByFieldID(fieldID); })
+      .Default([](Type) {
+        llvm_unreachable("unknown FIRRTL type");
+        return FIRRTLType();
+      });
+}
+
 std::pair<unsigned, bool> FIRRTLType::rootChildFieldID(unsigned fieldID,
                                                        unsigned index) {
   return TypeSwitch<FIRRTLType, std::pair<unsigned, bool>>(*this)
@@ -746,6 +758,12 @@ FIRRTLType BundleType::getElementType(StringRef name) {
   return element.hasValue() ? element.getValue().type : FIRRTLType();
 }
 
+FIRRTLType BundleType::getElementType(size_t index) {
+  assert(index < getNumElements() &&
+         "index must be less than number of fields in bundle");
+  return getElements()[index].type;
+}
+
 unsigned BundleType::getFieldID(unsigned index) {
   return getImpl()->fieldIDs[index];
 }
@@ -758,6 +776,15 @@ unsigned BundleType::getIndexForFieldID(unsigned fieldID) {
   return std::distance(fieldIDs.begin(), it);
 }
 
+FIRRTLType BundleType::getSubTypeByFieldID(unsigned fieldID) {
+  if (fieldID == 0)
+    return *this;
+  auto fieldIDs = getImpl()->fieldIDs;
+  auto it =
+      std::prev(std::upper_bound(fieldIDs.begin(), fieldIDs.end(), fieldID));
+  return getElementType(std::distance(fieldIDs.begin(), it));
+}
+
 unsigned BundleType::getMaxFieldID() { return getImpl()->maxFieldID; }
 
 std::pair<unsigned, bool> BundleType::rootChildFieldID(unsigned fieldID,
@@ -849,6 +876,12 @@ unsigned FVectorType::getIndexForFieldID(unsigned fieldID) {
   return (fieldID - 1) / (getElementType().getMaxFieldID() + 1);
 }
 
+FIRRTLType FVectorType::getSubTypeByFieldID(unsigned fieldID) {
+  if (fieldID == 0)
+    return *this;
+  return getElementType();
+}
+
 unsigned FVectorType::getMaxFieldID() {
   return getNumElements() * (getElementType().getMaxFieldID() + 1);
 }