Add fields location API

include/circt/Dialect/OM/OMOps.td

@@ -112,6 +112,28 @@ def ClassOp : OMClassLike<"class", [
     circt::om::ClassFieldsOp getFieldsOp() {
       return mlir::cast<ClassFieldsOp>(this->getBodyBlock()->getTerminator());
     }
+
+    // The addFields 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.
+    // This is required because MLIR's FusedLoc uses a "set" semantics where a
+    // single location is used to represent multiple fields with the same
+    // location.  The OM implementation uses the metadata attribute on FusedLoc
+    // 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>
+                   locs, mlir::ArrayRef<mlir::Value> values);
+
+    // Return the location for a field referenced by index in the fieldNames
+    // array attribute.  If the field has a location added by addFields 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
+    // assertion error
+    mlir::Location getFieldLocByIndex(size_t i);
   }];
 }
 

lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp

@@ -1151,7 +1151,7 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike,
   // Clone the property ops from the FIRRTL Class or Module to the OM Class.
   SmallVector<Operation *> opsToErase;
   OpBuilder builder = OpBuilder::atBlockBegin(classOp.getBodyBlock());
-  llvm::SmallVector<mlir::Attribute> fieldLocs;
+  llvm::SmallVector<mlir::Location> fieldLocs;
   llvm::SmallVector<mlir::Value> fieldValues;
   for (auto &op : moduleLike->getRegion(0).getOps()) {
     // Check if any operand is a property.
@@ -1173,7 +1173,7 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike,
     if (isa<PropAssignOp>(op) &&
         dyn_cast<BlockArgument>(cast<PropAssignOp>(op).getDest())) {
       // Store any output property assignments into fields op inputs.
-      fieldLocs.push_back(LocationAttr(op.getLoc()));
+      fieldLocs.push_back(op.getLoc());
       fieldValues.push_back(mapping.lookup(cast<PropAssignOp>(op).getSrc()));
     } else {
       // Clone the op over to the OM Class.
@@ -1190,14 +1190,11 @@ void LowerClassesPass::lowerClass(om::ClassOp classOp, FModuleLike moduleLike,
   if (hasContainingModule) {
     BlockArgument argumentValue = classBody->addArgument(
         getRtlPortsType(&getContext()), UnknownLoc::get(&getContext()));
-    fieldLocs.push_back(LocationAttr(argumentValue.getLoc()));
+    fieldLocs.push_back(argumentValue.getLoc());
     fieldValues.push_back(argumentValue);
   }
 
-  // TODO: Store set of locs?
-  // TODO: Document FusedLoc logic
-  builder.create<ClassFieldsOp>(
-      builder.getFusedLoc({}, builder.getArrayAttr(fieldLocs)), fieldValues);
+  classOp.addFields(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/Evaluator/Evaluator.cpp

@@ -16,7 +16,6 @@
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
-#include <iostream>
 
 #define DEBUG_TYPE "om-evaluator"
 
@@ -260,18 +259,10 @@ circt::om::Evaluator::evaluateObjectInstance(StringAttr className,
 
   auto fieldNames = cls.getFieldNames();
   auto operands = cls.getFieldsOp()->getOperands();
-  auto fieldsLoc = cls.getFieldsOp()->getLoc();
   for (size_t i = 0; i < fieldNames.size(); ++i) {
     auto name = fieldNames[i];
     auto value = operands[i];
-    Location fieldLoc = fieldsLoc;
-    if (auto locs = dyn_cast<FusedLoc>(fieldLoc)) {
-      // TODO: Document FusedLoc logic
-      assert(dyn_cast<ArrayAttr>(locs.getMetadata()));
-      assert(dyn_cast<LocationAttr>(cast<ArrayAttr>(locs.getMetadata())[i]));
-      fieldLoc =
-          Location(cast<LocationAttr>(cast<ArrayAttr>(locs.getMetadata())[i]));
-    }
+    auto fieldLoc = cls.getFieldLocByIndex(i);
     FailureOr<evaluator::EvaluatorValuePtr> result =
         evaluateValue(value, actualParams, fieldLoc);
     if (failed(result))

lib/Dialect/OM/OMOps.cpp

@@ -333,6 +333,38 @@ 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) {
+  // 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;
+  for (auto loc : locs) {
+    locAttrs.push_back(cast<Attribute>(LocationAttr(loc)));
+  }
+  // Also store the locations incase there's some other analysis that might
+  // be able to use the default FusedLoc representation.
+  builder.create<ClassFieldsOp>(
+      builder.getFusedLoc(locs, builder.getArrayAttr(locAttrs)), values);
+}
+
+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
+    ArrayAttr metadataArr = dyn_cast<ArrayAttr>(locs.getMetadata());
+    assert(metadataArr && "Expected fused loc to store metadata array");
+    assert(i < metadataArr.size() &&
+           "expected index to be less than array size");
+    LocationAttr locAttr = dyn_cast<LocationAttr>(metadataArr[i]);
+    assert(locAttr && "expected metadataArr entry to be location attribute");
+    loc = Location(locAttr);
+  }
+  return loc;
+}
+
 //===----------------------------------------------------------------------===//
 // ClassExternOp
 //===----------------------------------------------------------------------===//