[OM] Add a new API to update fields of a ClassOp (#7748)

Add a new API to update the fields of an existing `FieldsOp` in a class.
This also renames the `addFields` to ensure it is used only to add a new op.

include/circt/Dialect/OM/OMOps.td

@@ -121,7 +121,7 @@ def ClassOp : OMClassLike<"class", [
       return mlir::cast<ClassFieldsOp>(this->getBodyBlock()->getTerminator());
     }
 
-    // The addFields API encapsulates the logic used to represent field
+    // The addNewFieldsOp API encapsulates the logic used to represent field
     // locations under the hood.  Users should invoke this method rather
     // than construct the operation directly, otherwise logic to retrieve
     // the field location will break.
@@ -131,15 +131,22 @@ def ClassOp : OMClassLike<"class", [
     // to store the original array of locations, so that the specific location
     // of a field may be easily retrieved by index using the
     // `getFieldLocByIndex` API.
-    void addFields(mlir::OpBuilder &builder, mlir::ArrayRef<mlir::Location>
+    void addNewFieldsOp(mlir::OpBuilder &builder, mlir::ArrayRef<mlir::Location>
                    locs, mlir::ArrayRef<mlir::Value> values);
 
+    // Add additional fields to the class. This updates the ClassOp output type
+    //  with the additional field types and field names and adds the field
+    // values to the fieldsOp.
+    void updateFields(mlir::ArrayRef<mlir::Location> newLocations,
+                   mlir::ArrayRef<mlir::Value> newValues,
+                   mlir::ArrayRef<mlir::Attribute> newNames);
+
     // Return the location for a field referenced by index in the fieldNames
-    // array attribute.  If the field has a location added by addFields API,
+    // array attribute.  If the field has a location added by addNewFieldsOp API,
     // its location will be retrieved from the array of per field locations.
     // Otherwise, it will inherit the location of the class op Using this with
     // a ClassFieldsOp that has been constructed with a FusedLoc but not
-    // following the internal storage format of `addFields` will result in an
+    // following the internal storage format of `addNewFieldsOp` will result in an
     // assertion error
     mlir::Location getFieldLocByIndex(size_t i);
   }];

lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp

@@ -1201,7 +1201,7 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike,
     fieldValues.push_back(argumentValue);
   }
 
-  classOp.addFields(builder, fieldLocs, fieldValues);
+  classOp.addNewFieldsOp(builder, fieldLocs, fieldValues);
 
   // If the module-like is a Class, it will be completely erased later.
   // Otherwise, erase just the property ports and ops.

lib/Dialect/OM/OMOps.cpp

@@ -349,9 +349,64 @@ void circt::om::ClassOp::replaceFieldTypes(AttrTypeReplacer replacer) {
   replaceClassLikeFieldTypes(*this, replacer);
 }
 
-void circt::om::ClassOp::addFields(mlir::OpBuilder &builder,
-                                   mlir::ArrayRef<Location> locs,
-                                   mlir::ArrayRef<Value> values) {
+void circt::om::ClassOp::updateFields(
+    mlir::ArrayRef<mlir::Location> newLocations,
+    mlir::ArrayRef<mlir::Value> newValues,
+    mlir::ArrayRef<mlir::Attribute> newNames) {
+
+  auto fieldsOp = getFieldsOp();
+  assert(fieldsOp && "The fields op should exist");
+  // Get field names.
+  SmallVector<Attribute> names(getFieldNamesAttr().getAsRange<StringAttr>());
+  // Get the field types.
+  SmallVector<NamedAttribute> fieldTypes(getFieldTypesAttr().getValue());
+  // Get the field values.
+  SmallVector<Value> fieldVals(fieldsOp.getFields());
+  // Get the field locations.
+  Location fieldOpLoc = fieldsOp->getLoc();
+
+  // Extract the locations per field.
+  SmallVector<Location> locations;
+  if (auto fl = dyn_cast<FusedLoc>(fieldOpLoc)) {
+    auto metadataArr = dyn_cast<ArrayAttr>(fl.getMetadata());
+    assert(metadataArr && "Expected the metadata for the fused location");
+    auto r = metadataArr.getAsRange<LocationAttr>();
+    locations.append(r.begin(), r.end());
+  } else {
+    // Assume same loc for every field.
+    locations.append(names.size(), fieldOpLoc);
+  }
+
+  // Append the new names, locations and values.
+  names.append(newNames.begin(), newNames.end());
+  locations.append(newLocations.begin(), newLocations.end());
+  fieldVals.append(newValues.begin(), newValues.end());
+
+  // Construct the new field types from values and names.
+  for (auto [v, n] : llvm::zip(newValues, newNames))
+    fieldTypes.emplace_back(
+        NamedAttribute(llvm::cast<StringAttr>(n), TypeAttr::get(v.getType())));
+
+  // Keep the locations as array on the metadata.
+  SmallVector<Attribute> locationsAttr;
+  llvm::for_each(locations, [&](Location &l) {
+    locationsAttr.push_back(cast<Attribute>(l));
+  });
+
+  ImplicitLocOpBuilder builder(getLoc(), *this);
+  // Update the field names attribute.
+  setFieldNamesAttr(builder.getArrayAttr(names));
+  // Update the fields type attribute.
+  setFieldTypesAttr(builder.getDictionaryAttr(fieldTypes));
+  fieldsOp.getFieldsMutable().assign(fieldVals);
+  // Update the location.
+  fieldsOp->setLoc(builder.getFusedLoc(
+      locations, ArrayAttr::get(getContext(), locationsAttr)));
+}
+
+void circt::om::ClassOp::addNewFieldsOp(mlir::OpBuilder &builder,
+                                        mlir::ArrayRef<Location> locs,
+                                        mlir::ArrayRef<Value> values) {
   // Store the original locations as a metadata array so that unique locations
   // are preserved as a mapping from field index to location
   mlir::SmallVector<Attribute> locAttrs;
@@ -368,8 +423,8 @@ mlir::Location circt::om::ClassOp::getFieldLocByIndex(size_t i) {
   Location loc = this->getFieldsOp()->getLoc();
   if (auto locs = dyn_cast<FusedLoc>(loc)) {
     // Because it's possible for a user to construct a fields op directly and
-    // place a FusedLoc that doersn't follow the storage format of addFields, we
-    // assert the information has been stored appropriately
+    // place a FusedLoc that doersn't follow the storage format of
+    // addNewFieldsOp, we assert the information has been stored appropriately
     ArrayAttr metadataArr = dyn_cast<ArrayAttr>(locs.getMetadata());
     assert(metadataArr && "Expected fused loc to store metadata array");
     assert(i < metadataArr.size() &&