[FIRRTL][InferReset] Keep instance graph properly updated (#1912)

In an earlier change, I made InferResets always mark the InstanceGraph
analysis as preserved, under the assumption that the pass does not
modify the instance graph.  This turns out to be incorrect as it can add
arguments to a module and must recreate the instance ops to reflect
that.

This change adds a function `InstanceGraph::replaceInstance(old, new)`
which allows the InstanceGraph to stay updated by the pass. This is the
first function added to the InstanceGraph to allow updating the cached
analysis.  This function is slightly more efficient then adding a
separate `addInstance(new)` and removeInstance(old)` functions.

include/circt/Dialect/FIRRTL/InstanceGraph.h

@@ -41,6 +41,8 @@ public:
   InstanceGraphNode *getTarget() const { return target; }
 
 private:
+  friend class InstanceGraph;
+
   /// The InstanceOp that this is tracking.
   InstanceOp instance;
 
@@ -169,6 +171,18 @@ public:
   iterator begin() { return nodes.begin(); }
   iterator end() { return nodes.end(); }
 
+  //===-------------------------------------------------------------------------
+  // Methods to keep an InstanceGraph up to date.
+  //
+  // These methods are not thread safe.  Make sure that modifications are
+  // properly synchronized or performed in a serial context.  When the
+  // InstanceGraph is used as an analysis, this is only safe when the pass is
+  // on a CircuitOp.
+
+  // Replaces an instance of a module with another instance. The target module
+  // of both InstanceOps must be the same.
+  void replaceInstance(InstanceOp inst, InstanceOp newInst);
+
 private:
   /// Get the node corresponding to the module.  If the node has does not exist
   /// yet, it will be created.

lib/Dialect/FIRRTL/InstanceGraph.cpp

@@ -94,6 +94,22 @@ Operation *InstanceGraph::getReferencedModule(InstanceOp op) {
   return lookup(op.moduleName())->getModule();
 }
 
+void InstanceGraph::replaceInstance(InstanceOp inst, InstanceOp newInst) {
+  assert(inst.moduleName() == newInst.moduleName() &&
+         "Both instances must be targeting the same module");
+
+  // Find the instance record of this instance.
+  auto *node = lookup(inst.moduleName());
+  auto it = llvm::find_if(node->uses(), [&](InstanceRecord *record) {
+    return record->getInstance() == inst;
+  });
+  assert(it != node->uses_end() && "Instance of module not recorded in graph");
+
+  // We can just replace the instance op in the InstanceRecord without updating
+  // any instance lists.
+  (*it)->instance = newInst;
+}
+
 ArrayRef<InstancePath> InstancePathCache::getAbsolutePaths(Operation *op) {
   assert((isa<FModuleOp, FExtModuleOp>(op))); // extra parens makes parser smile
 

lib/Dialect/FIRRTL/Transforms/InferResets.cpp

@@ -1459,6 +1459,7 @@ void InferResetsPass::implementAsyncReset(Operation *op, FModuleOp module,
 
       // Update the uses over to the new instance and drop the old instance.
       instOp.replaceAllUsesWith(newInstOp.getResults().drop_front());
+      instanceGraph->replaceInstance(instOp, newInstOp);
       instOp->erase();
       instOp = newInstOp;
     } else if (domain.existingPort.hasValue()) {