[FIRRTL] Update LowerClasses to alter what PathInfo stores. (#7522)

PathInfo stores a pointer to an Operation, which was problematic
because that Operation may be deleted in updateInstanceInModule. The
test case added in this patch would lead to a use-after-free.

This pointer was only really used for a few things, which can be
handled differently to avoid needing to consider Operation lifetimes.

One use was the operator bool implementation, to check if a PathInfo is
empty. In the one place this was used, an equivalent check is to query
the PathInfoTable, and check if the key was not in the table.

Another use was adding the Operation's location in error messages and
notes. We can safely store a Location directly for these messages.

The final use was to do an isa check while determining the target
kind. This is where the use in the use-after-free would manifest. For
this, we do the isa check early, and store the result in a bool.

In summary, we are able to simplify the data in PathInfo in order to
avoid hanging on to an Operation pointer and needing to worry about
its lifetime.

Author: mikeurbach

lib/Dialect/FIRRTL/Transforms/LowerClasses.cpp

@@ -81,17 +81,18 @@ static bool shouldCreateClassImpl(igraph::InstanceGraphNode *node) {
 /// to the targeted operation.
 struct PathInfo {
   PathInfo() = default;
-  PathInfo(Operation *op, FlatSymbolRefAttr symRef,
+  PathInfo(Location loc, bool canBeInstanceTarget, FlatSymbolRefAttr symRef,
            StringAttr altBasePathModule)
-      : op(op), symRef(symRef), altBasePathModule(altBasePathModule) {
-    assert(op && "op must not be null");
+      : loc(loc), canBeInstanceTarget(canBeInstanceTarget), symRef(symRef),
+        altBasePathModule(altBasePathModule) {
     assert(symRef && "symRef must not be null");
   }
 
-  operator bool() const { return op != nullptr; }
+  /// The Location of the hardware component targeted by this path.
+  std::optional<Location> loc = std::nullopt;
 
-  /// The hardware component targeted by this path.
-  Operation *op = nullptr;
+  /// Flag to indicate if the hardware component can be targeted as an instance.
+  bool canBeInstanceTarget = false;
 
   /// A reference to the hierarchical path targeting the op.
   FlatSymbolRefAttr symRef = nullptr;
@@ -584,17 +585,23 @@ LogicalResult PathTracker::updatePathInfoTable(PathInfoTable &pathInfoTable,
     auto [it, inserted] = pathInfoTable.table.try_emplace(entry.id);
     auto &pathInfo = it->second;
     if (!inserted) {
-      auto diag =
-          emitError(pathInfo.op->getLoc(), "duplicate identifier found");
+      assert(pathInfo.loc.has_value() && "all PathInfo should have a Location");
+      auto diag = emitError(pathInfo.loc.value(), "duplicate identifier found");
       diag.attachNote(entry.op->getLoc()) << "other identifier here";
       return failure();
     }
 
-    if (entry.pathAttr)
-      pathInfo = {entry.op, cache.getRefFor(entry.pathAttr),
-                  entry.altBasePathModule};
-    else
-      pathInfo.op = entry.op;
+    // Check if the op is targetable by an instance target. The op pointer may
+    // be invalidated later, so this is the last time we want to access it here.
+    bool canBeInstanceTarget = isa<InstanceOp, FModuleLike>(entry.op);
+
+    if (entry.pathAttr) {
+      pathInfo = {entry.op->getLoc(), canBeInstanceTarget,
+                  cache.getRefFor(entry.pathAttr), entry.altBasePathModule};
+    } else {
+      pathInfo.loc = entry.op->getLoc();
+      pathInfo.canBeInstanceTarget = canBeInstanceTarget;
+    }
   }
   return success();
 }
@@ -1453,14 +1460,14 @@ struct PathOpConversion : public OpConversionPattern<firrtl::PathOp> {
                   ConversionPatternRewriter &rewriter) const override {
     auto *context = op->getContext();
     auto pathType = om::PathType::get(context);
-    auto pathInfo = pathInfoTable.table.lookup(op.getTarget());
+    auto pathInfoIt = pathInfoTable.table.find(op.getTarget());
 
     // The 0'th argument is the base path by default.
     auto basePath = op->getBlock()->getArgument(0);
 
     // If the target was optimized away, then replace the path operation with
     // a deleted path.
-    if (!pathInfo) {
+    if (pathInfoIt == pathInfoTable.table.end()) {
       if (op.getTargetKind() == firrtl::TargetKind::DontTouch)
         return emitError(op.getLoc(), "DontTouch target was deleted");
       if (op.getTargetKind() == firrtl::TargetKind::Instance)
@@ -1469,6 +1476,7 @@ struct PathOpConversion : public OpConversionPattern<firrtl::PathOp> {
       return success();
     }
 
+    auto pathInfo = pathInfoIt->second;
     auto symbol = pathInfo.symRef;
 
     // Convert the target kind to an OMIR target.  Member references are updated
@@ -1482,15 +1490,15 @@ struct PathOpConversion : public OpConversionPattern<firrtl::PathOp> {
       targetKind = om::TargetKind::Reference;
       break;
     case firrtl::TargetKind::Instance:
-      if (!isa<InstanceOp, FModuleLike>(pathInfo.op))
+      if (!pathInfo.canBeInstanceTarget)
         return emitError(op.getLoc(), "invalid target for instance path")
-                   .attachNote(pathInfo.op->getLoc())
+                   .attachNote(pathInfo.loc)
                << "target not instance or module";
       targetKind = om::TargetKind::Instance;
       break;
     case firrtl::TargetKind::MemberInstance:
     case firrtl::TargetKind::MemberReference:
-      if (isa<InstanceOp, FModuleLike>(pathInfo.op))
+      if (pathInfo.canBeInstanceTarget)
         targetKind = om::TargetKind::MemberInstance;
       else
         targetKind = om::TargetKind::MemberReference;

test/Dialect/FIRRTL/lower-classes.mlir

@@ -458,3 +458,16 @@ firrtl.circuit "OwningModulePrefix" {
     firrtl.path reference distinct[0]<>
   }
 }
+
+// CHECK-LABEL: firrtl.circuit "PathTargetReplaced"
+firrtl.circuit "PathTargetReplaced" {
+  // CHECK: hw.hierpath private [[NLA:@.+]] [@PathTargetReplaced::[[SYM:@.+]]]
+  firrtl.module @PathTargetReplaced() {
+    // CHECK: firrtl.instance replaced sym [[SYM]]
+    firrtl.instance replaced {annotations = [{class = "circt.tracker", id = distinct[0]<>}]} @WillBeReplaced(out output: !firrtl.integer)
+    // CHECK: om.path_create instance %basepath [[NLA]]
+    %path = firrtl.path instance distinct[0]<>
+  }
+  firrtl.module private @WillBeReplaced(out %output: !firrtl.integer) {
+  }
+}