[FIRRTL] Add utilities for updating NLA. NFC. (#2447)

Add helper utilities for updating the `FIRRTL::NonLocalAnchor` op.

include/circt/Dialect/FIRRTL/FIRRTLStructure.td

@@ -146,4 +146,26 @@ def NonLocalAnchor : FIRRTLOp<"nla",
   let arguments = (ins SymbolNameAttr:$sym_name, NameRefArrayAttr:$namepath);
   let results = (outs);
   let assemblyFormat = [{ $sym_name $namepath attr-dict}];
+  let extraClassDeclaration = [{
+    /// Drop the module from the namepath. If its a InnerNameRef, then drop
+    /// the Module-Instance pair, else drop the final module from the namepath.
+    /// Return true if any update is made.
+    bool dropModule(StringAttr moduleToDrop);
+
+    /// Inline the module in the namepath.
+    /// Update the symbol name for the inlined module instance, by prepending
+    /// the symbol name of the instance at which the inling was done.
+    /// Return true if any update is made.
+     bool inlineModule(StringAttr moduleToDrop);
+
+    /// Replace the oldMod module with newMod module in the namepath of the NLA.
+    /// Return true if any update is made.
+    bool updateModule(StringAttr oldMod, StringAttr newMod);
+
+    /// Truncate the namepath for this NLA, at atMod module.
+    /// If includeMod is false, drop atMod and beyond, else include it and drop
+    /// everything after it.
+    /// Return true if any update is made.
+    bool truncateAtModule(StringAttr atMod, bool includeMod = true);
+    }];
 }

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -3025,6 +3025,112 @@ static void printVerifAttrs(OpAsmPrinter &p, Operation *op,
   printElideEmptyName(p, op, attr, {"message"});
 }
 
+//===----------------------------------------------------------------------===//
+// NonLocalAnchor helpers.
+//===----------------------------------------------------------------------===//
+
+bool NonLocalAnchor::dropModule(StringAttr moduleToDrop) {
+  SmallVector<Attribute, 4> newPath;
+  bool updateMade = false;
+  for (auto nameRef : namepath()) {
+    // nameRef is either an InnerRefAttr or a FlatSymbolRefAttr.
+    if (auto ref = nameRef.dyn_cast<hw::InnerRefAttr>()) {
+      if (ref.getModule() == moduleToDrop)
+        updateMade = true;
+      else
+        newPath.push_back(ref);
+    } else {
+      if (nameRef.cast<FlatSymbolRefAttr>().getAttr() == moduleToDrop)
+        updateMade = true;
+      else
+        newPath.push_back(nameRef);
+    }
+  }
+  if (updateMade)
+    namepathAttr(ArrayAttr::get(getContext(), newPath));
+  return updateMade;
+}
+
+bool NonLocalAnchor::inlineModule(StringAttr moduleToDrop) {
+  SmallVector<Attribute, 4> newPath;
+  bool updateMade = false;
+  StringRef inlinedInstanceName = "";
+  for (auto nameRef : namepath()) {
+    // nameRef is either an InnerRefAttr or a FlatSymbolRefAttr.
+    if (auto ref = nameRef.dyn_cast<hw::InnerRefAttr>()) {
+      if (ref.getModule() == moduleToDrop) {
+        inlinedInstanceName = ref.getName().getValue();
+        updateMade = true;
+      } else if (!inlinedInstanceName.empty()) {
+        newPath.push_back(hw::InnerRefAttr::get(
+            getContext(), ref.getModule(),
+            StringAttr::get(getContext(), inlinedInstanceName + "_" +
+                                              ref.getName().getValue())));
+        inlinedInstanceName = "";
+      } else
+        newPath.push_back(ref);
+    } else {
+      if (nameRef.cast<FlatSymbolRefAttr>().getAttr() == moduleToDrop)
+        updateMade = true;
+      else
+        newPath.push_back(nameRef);
+    }
+  }
+  if (updateMade)
+    namepathAttr(ArrayAttr::get(getContext(), newPath));
+  return updateMade;
+}
+
+bool NonLocalAnchor::updateModule(StringAttr oldMod, StringAttr newMod) {
+  SmallVector<Attribute, 4> newPath;
+  bool updateMade = false;
+  for (auto nameRef : namepath()) {
+    // nameRef is either an InnerRefAttr or a FlatSymbolRefAttr.
+    if (auto ref = nameRef.dyn_cast<hw::InnerRefAttr>()) {
+      if (ref.getModule() == oldMod) {
+        newPath.push_back(hw::InnerRefAttr::get(newMod, ref.getName()));
+        updateMade = true;
+      } else
+        newPath.push_back(ref);
+    } else {
+      if (nameRef.cast<FlatSymbolRefAttr>().getAttr() == oldMod) {
+        newPath.push_back(FlatSymbolRefAttr::get(newMod));
+        updateMade = true;
+      } else
+        newPath.push_back(nameRef);
+    }
+  }
+  if (updateMade)
+    namepathAttr(ArrayAttr::get(getContext(), newPath));
+  return updateMade;
+}
+
+bool NonLocalAnchor::truncateAtModule(StringAttr atMod, bool includeMod) {
+  SmallVector<Attribute, 4> newPath;
+  bool updateMade = false;
+  for (auto nameRef : namepath()) {
+    // nameRef is either an InnerRefAttr or a FlatSymbolRefAttr.
+    if (auto ref = nameRef.dyn_cast<hw::InnerRefAttr>()) {
+      if (ref.getModule() == atMod) {
+        updateMade = true;
+        if (includeMod)
+          newPath.push_back(ref);
+      } else
+        newPath.push_back(ref);
+    } else {
+      if (nameRef.cast<FlatSymbolRefAttr>().getAttr() == atMod && !includeMod)
+        updateMade = true;
+      else
+        newPath.push_back(nameRef);
+    }
+    if (updateMade)
+      break;
+  }
+  if (updateMade)
+    namepathAttr(ArrayAttr::get(getContext(), newPath));
+  return updateMade;
+}
+
 //===----------------------------------------------------------------------===//
 // TblGen Generated Logic.
 //===----------------------------------------------------------------------===//

lib/Dialect/FIRRTL/Transforms/PrefixModules.cpp

@@ -168,28 +168,9 @@ void PrefixModulesPass::renameModuleBody(std::string prefix, FModuleOp module) {
               instanceOp.emitError("cannot find NonLocalAnchor :" + nlaName);
             else {
               auto nlaOp = dyn_cast<NonLocalAnchor>(f->second);
-              // Iterate over the modules of the NonLocalAnchor op, and update
-              // it.
-              SmallVector<Attribute, 4> newMods;
-              for (auto nameRef : nlaOp.namepath()) {
-                // nameRef is either an InnerRefAttr or a FlatSymbolRefAttr.
-                if (auto oldMod = nameRef.dyn_cast<hw::InnerRefAttr>()) {
-                  if (instanceOp.moduleNameAttr().getAttr() ==
-                      oldMod.getModule())
-                    newMods.push_back(hw::InnerRefAttr::get(
-                        StringAttr::get(context, newTarget), oldMod.getName()));
-                  else
-                    newMods.push_back(oldMod);
-                } else {
-                  if (instanceOp.moduleNameAttr() ==
-                      nameRef.cast<FlatSymbolRefAttr>())
-                    newMods.push_back(
-                        FlatSymbolRefAttr::get(context, newTarget));
-                  else
-                    newMods.push_back(nameRef);
-                }
-              }
-              nlaOp->setAttr("namepath", ArrayAttr::get(context, newMods));
+              StringAttr oldModName = instanceOp.moduleNameAttr().getAttr();
+              nlaOp.updateModule(oldModName,
+                                 StringAttr::get(context, newTarget));
             }
           }
       }