[FIRRTL] Fix port insertion/deletion domain bugs

Fix bugs in the logic for port insertion and deletion.

This was factored out of #9106.  I'm trying to fixing bugs in _other_
passes, without having to land the full domain inference pass just yet and
this was a prerequisite for some of those other fixes.

Co-authored-by: Robert Young <[email protected]>
Signed-off-by: Schuyler Eldridge <[email protected]>

Author: seldridge

lib/Dialect/FIRRTL/FIRRTLOps.cpp

@@ -801,6 +801,15 @@ static void insertPorts(FModuleLike op,
   unsigned oldNumArgs = op.getNumPorts();
   unsigned newNumArgs = oldNumArgs + ports.size();
 
+  // Build a map from old port indices to new indices.
+  SmallVector<unsigned> indexMap(oldNumArgs);
+  size_t inserted = 0;
+  for (size_t i = 0; i < oldNumArgs; ++i) {
+    while (inserted < ports.size() && ports[inserted].first == i)
+      ++inserted;
+    indexMap[i] = i + inserted;
+  }
+
   // Add direction markers and names for new ports.
   auto existingDirections = op.getPortDirectionsAttr();
   ArrayRef<Attribute> existingNames = op.getPortNames();
@@ -822,12 +831,9 @@ static void insertPorts(FModuleLike op,
   newSyms.reserve(newNumArgs);
   newLocs.reserve(newNumArgs);
 
-  SmallVector<unsigned> indexMap(oldNumArgs);
-
   auto emptyArray = ArrayAttr::get(op.getContext(), {});
 
   unsigned oldIdx = 0;
-  unsigned inserted = 0;
   auto migrateOldPorts = [&](unsigned untilOldIdx) {
     while (oldIdx < oldNumArgs && oldIdx < untilOldIdx) {
       newDirections.push_back(existingDirections[oldIdx]);
@@ -838,8 +844,6 @@ static void insertPorts(FModuleLike op,
       newAnnos.push_back(op.getAnnotationsAttrForPort(oldIdx));
       newSyms.push_back(op.getPortSymbolAttr(oldIdx));
       newLocs.push_back(existingLocs[oldIdx]);
-
-      indexMap[oldIdx] = oldIdx + inserted;
       ++oldIdx;
     }
   };
@@ -905,24 +909,13 @@ static ArrayAttr fixDomainInfoDeletions(MLIRContext *context,
   if (supportsEmptyAttr && domainInfoAttr.empty())
     return domainInfoAttr;
 
-  // Compute an array where each entry is the number of ports before that
-  // index that have been deleted.  This is used to update domain information.
-  SmallVector<unsigned> numDeleted;
-  numDeleted.resize(portIndices.size());
-  size_t deletionIndex = portIndices.find_first();
+  // Build a map from old port indices to new indices.
+  SmallVector<unsigned> indexMap(portIndices.size());
+  size_t deleted = 0;
   for (size_t i = 0, e = portIndices.size(); i != e; ++i) {
-    if (i == deletionIndex) {
-      if (i == 0)
-        numDeleted[i] = 1;
-      else
-        numDeleted[i] = numDeleted[i - 1] + 1;
-      deletionIndex = portIndices.find_next(i);
-      continue;
-    }
-    if (i == 0)
-      numDeleted[i] = 0;
-    else
-      numDeleted[i] = numDeleted[i - 1];
+    indexMap[i] = i - deleted;
+    if (portIndices[i])
+      ++deleted;
   }
 
   // Return a cached empty ArrayAttr.
@@ -960,7 +953,7 @@ static ArrayAttr fixDomainInfoDeletions(MLIRContext *context,
       if (portIndices.test(oldIdx))
         continue;
       // If the new index is the same, do nothing.
-      auto newIdx = oldIdx - numDeleted[oldIdx];
+      auto newIdx = indexMap[oldIdx];
       if (oldIdx == newIdx) {
         newDomainInfo.push_back(attr);
         continue;