[LowerToHW] Fix output port index mapping (#6530)

The logic that was moving the symbol from the port to the wire was
 incorrectly assuming that all the output ports are ordered after the input
  ports, this resulted in the port index mapping being incorrect.
This was causing the incorrect port symbol being moved to a temporary wire.
This commit fixes the port index calculation logic.

lib/Conversion/FIRRTLToHW/LowerToHW.cpp

@@ -1341,30 +1341,33 @@ LogicalResult FIRRTLModuleLowering::lowerModulePortsAndMoveBody(
   bodyBuilder.setInsertionPoint(cursor);
 
   // Insert argument casts, and re-vector users in the old body to use them.
-  SmallVector<PortInfo> ports = oldModule.getPorts();
-  assert(oldModule.getBody().getNumArguments() == ports.size() &&
+  SmallVector<PortInfo> firrtlPorts = oldModule.getPorts();
+  SmallVector<hw::PortInfo> hwPorts = newModule.getPortList();
+  assert(oldModule.getBody().getNumArguments() == firrtlPorts.size() &&
          "port count mismatch");
 
-  size_t nextNewArg = 0;
-  size_t firrtlArg = 0;
   SmallVector<Value, 4> outputs;
 
   // This is the terminator in the new module.
   auto outputOp = newModule.getBodyBlock()->getTerminator();
   ImplicitLocOpBuilder outputBuilder(oldModule.getLoc(), outputOp);
 
-  for (auto &port : ports) {
+  unsigned nextHWInputArg = 0;
+  int hwPortIndex = -1;
+  for (auto [firrtlPortIndex, port] : llvm::enumerate(firrtlPorts)) {
     // Inputs and outputs are both modeled as arguments in the FIRRTL level.
-    auto oldArg = oldModule.getBody().getArgument(firrtlArg++);
+    auto oldArg = oldModule.getBody().getArgument(firrtlPortIndex);
 
     bool isZeroWidth =
         type_isa<FIRRTLBaseType>(port.type) &&
         type_cast<FIRRTLBaseType>(port.type).getBitWidthOrSentinel() == 0;
+    if (!isZeroWidth)
+      ++hwPortIndex;
 
     if (!port.isOutput() && !isZeroWidth) {
       // Inputs and InOuts are modeled as arguments in the result, so we can
       // just map them over.  We model zero bit outputs as inouts.
-      Value newArg = newModule.getBody().getArgument(nextNewArg++);
+      Value newArg = newModule.getBody().getArgument(nextHWInputArg++);
 
       // Cast the argument to the old type, reintroducing sign information in
       // the hw.module body.
@@ -1407,13 +1410,12 @@ LogicalResult FIRRTLModuleLowering::lowerModulePortsAndMoveBody(
     if (!resultHWType.isInteger(0)) {
       auto output =
           castFromFIRRTLType(newArg.getResult(), resultHWType, outputBuilder);
-      auto idx = newModule.getNumInputPorts() + outputs.size();
       outputs.push_back(output);
 
       // If output port has symbol, move it to this wire.
-      if (auto sym = newModule.getPortList()[idx].getSym()) {
+      if (auto sym = hwPorts[hwPortIndex].getSym()) {
         newArg.setInnerSymAttr(sym);
-        newModule.setPortSymbolAttr(idx, {});
+        newModule.setPortSymbolAttr(hwPortIndex, {});
       }
     }
   }

test/Conversion/FIRRTLToHW/lower-to-hw.mlir

@@ -1581,3 +1581,22 @@ firrtl.circuit "ZeroWidthForeignOperand" {
     dbg.variable "v2", %a : !firrtl.uint<0>
   }
 }
+
+// -----
+
+// Check correct symbol mapping from output port to wire happens.
+firrtl.circuit "PortSym" {
+  firrtl.extmodule private @Blackbox(out bar: !firrtl.uint<1>)
+  // CHECK-LABEL: module @PortSym(
+  // CHECK-SAME: out a : i1 {hw.exportPort = #hw<innerSym@out_a_m>}
+  // CHECK-SAME: out out : i5, in %c : i1)
+  firrtl.module @PortSym(out %a: !firrtl.uint<1> sym @out_a_m, in %b: !firrtl.uint<0>, out %out: !firrtl.uint<5> sym @out_sym, in %c: !firrtl.uint<1>) attributes {convention = #firrtl<convention scalarized>} {
+    // CHECK: %[[OUT:.+]] = hw.wire %{{.+}} sym @out_sym
+    %c1_ui1 = firrtl.constant 1 : !firrtl.uint<1>
+    %c1_ui5 = firrtl.constant 1 : !firrtl.uint<5>
+    firrtl.strictconnect %out, %c1_ui5 : !firrtl.uint<5>
+    %e_a = firrtl.instance sub1 @Blackbox(out bar: !firrtl.uint<1>)
+    firrtl.strictconnect %a, %e_a : !firrtl.uint<1>
+    %0 = firrtl.eq %out, %c1_ui5 : (!firrtl.uint<5>, !firrtl.uint<5>) -> !firrtl.uint<1>
+  }
+}