🐛 Fix one-time-use of variables for the ctrl modifier

mlir/include/mlir/Dialect/QCO/IR/QCOOps.td

@@ -1025,10 +1025,10 @@ def YieldOp : QCOOp<"yield", traits = [Terminator]> {
 
         Example:
         ```mlir
-        %ctrl_q_out, %tgt_q_out = qco.ctrl(%ctrl_q_in) %tgt_q_in {
-            %tgt_q_res = qco.h %tgt_q_in : !qco.qubit -> !qco.qubit
-            qco.yield %tgt_q_res : !qco.qubit
-        } : ({!qco.qubit}, {!qco.qubit}) -> ({!qco.qubit}, {!qco.qubit})
+        %res_ctrl, %res_tgt:2 = qco.ctrl(%ctrl) targets(%a0 = %q0, %a1 = %q1) {
+            %a0_1, %a1_1 = qco.swap %a0, %a1 : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit
+            qco.yield %a0_1, %a1_1
+        } : ({!qco.qubit}, {!qco.qubit, !qco.qubit}) -> ({!qco.qubit}, {!qco.qubit, !qco.qubit})
         ```
     }];
 
@@ -1057,10 +1057,10 @@ def CtrlOp : QCOOp<"ctrl", traits =
 
         Example:
         ```mlir
-        %ctrl_q_out, %tgt_q_out = qco.ctrl(%ctrl_q_in) %tgt_q_in {
-            %tgt_q_res = qco.h %tgt_q_in : !qco.qubit -> !qco.qubit
-            qco.yield %tgt_q_res : !qco.qubit
-        } : ({!qco.qubit}, {!qco.qubit}) -> ({!qco.qubit}, {!qco.qubit})
+        %res_ctrl, %res_tgt:2 = qco.ctrl(%ctrl) targets(%a0 = %q0, %a1 = %q1) {
+            %a0_1, %a1_1 = qco.swap %a0, %a1 : !qco.qubit, !qco.qubit -> !qco.qubit, !qco.qubit
+            qco.yield %a0_1, %a1_1
+        } : ({!qco.qubit}, {!qco.qubit, !qco.qubit}) -> ({!qco.qubit}, {!qco.qubit, !qco.qubit})
         ```
     }];
 
@@ -1069,8 +1069,10 @@ def CtrlOp : QCOOp<"ctrl", traits =
     let results = (outs Variadic<QubitType>:$controls_out, Variadic<QubitType>:$targets_out);
     let regions = (region SizedRegion<1>:$region);
     let assemblyFormat = [{
-        `(` $controls_in `)` $targets_in
-        $region attr-dict `:`
+        `(` $controls_in `)`
+        `targets`
+        custom<TargetAliasing>($region, $targets_in)
+        attr-dict `:`
         `(` `{` type($controls_in) `}` ( `,` `{` type($targets_in)^ `}` )? `)`
         `->`
         `(` `{` type($controls_out) `}` ( `,` `{` type($targets_out)^ `}` )? `)`

mlir/lib/Conversion/QCOToQC/QCOToQC.cpp

@@ -774,12 +774,32 @@ struct ConvertQCOCtrlOp final : OpConversionPattern<qco::CtrlOp> {
     const auto& qcControls = adaptor.getControlsIn();
 
     // Create qc.ctrl operation
-    auto qcoOp = rewriter.create<qc::CtrlOp>(op.getLoc(), qcControls);
+    auto qcOp = rewriter.create<qc::CtrlOp>(op.getLoc(), qcControls);
 
     // Clone body region from QCO to QC
-    auto& dstRegion = qcoOp.getRegion();
+    auto& dstRegion = qcOp.getRegion();
     rewriter.cloneRegionBefore(op.getRegion(), dstRegion, dstRegion.end());
 
+    // Remove all block arguments in the cloned region
+    rewriter.modifyOpInPlace(qcOp, [&] {
+      auto& entryBlock = dstRegion.front();
+      auto numArgs = entryBlock.getNumArguments();
+
+      // 1. Replace uses (Must be done BEFORE erasing)
+      // We iterate 0..N using indices since the block args are still stable
+      // here.
+      for (auto i = 0UL; i < numArgs; ++i) {
+        entryBlock.getArgument(i).replaceAllUsesWith(adaptor.getTargetsIn()[i]);
+      }
+
+      // 2. Erase all block arguments
+      // Now that they have no uses, we can safely wipe them.
+      // We use a bulk erase for efficiency (start index 0, count N).
+      if (numArgs > 0) {
+        entryBlock.eraseArguments(0, numArgs);
+      }
+    });
+
     // Replace the output qubits with the same QC references
     rewriter.replaceOp(op, adaptor.getOperands());
 

mlir/lib/Conversion/QCToQCO/QCToQCO.cpp

@@ -1124,12 +1124,27 @@ struct ConvertQCCtrlOp final : StatefulOpConversionPattern<qc::CtrlOp> {
 
     // Update modifier information
     state.inCtrlOp++;
-    state.targetsIn.try_emplace(state.inCtrlOp, qcoTargets);
 
     // Clone body region from QC to QCO
     auto& dstRegion = qcoOp.getRegion();
     rewriter.cloneRegionBefore(op.getRegion(), dstRegion, dstRegion.end());
 
+    // Create block arguments for target qubits and store them in
+    // `state.targetsIn`.
+    auto& entryBlock = dstRegion.front();
+    SmallVector<Value> qcoTargetAliases;
+    qcoTargetAliases.reserve(numTargets);
+
+    rewriter.modifyOpInPlace(qcoOp, [&] {
+      for (auto i = 0UL; i < qcoOp.getNumTargets(); i++) {
+        auto originalTarget = qcoOp.getInputTarget(i);
+        auto arg =
+            entryBlock.addArgument(originalTarget.getType(), op.getLoc());
+        qcoTargetAliases.push_back(arg);
+      }
+    });
+    state.targetsIn.try_emplace(state.inCtrlOp, qcoTargetAliases);
+
     rewriter.eraseOp(op);
     return success();
   }

mlir/lib/Dialect/QCO/IR/QCOOps.cpp

@@ -9,6 +9,11 @@
  */
 
 #include "mlir/Dialect/QCO/IR/QCODialect.h" // IWYU pragma: associated
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/Region.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/SmallVector.h"
 
 // The following headers are needed for some template instantiations.
 // IWYU pragma: begin_keep
@@ -21,6 +26,85 @@
 using namespace mlir;
 using namespace mlir::qco;
 
+//===----------------------------------------------------------------------===//
+// Custom Parsers
+//===----------------------------------------------------------------------===//
+
+namespace {
+ParseResult
+parseTargetAliasing(OpAsmParser& parser, Region& region,
+                    SmallVectorImpl<OpAsmParser::UnresolvedOperand>& operands) {
+  // 1. Parse the opening parenthesis
+  if (parser.parseLParen()) {
+    return failure();
+  }
+
+  // Temporary storage for block arguments we are about to create
+  SmallVector<OpAsmParser::Argument> blockArgs;
+
+  // 2. Prepare to parse the list
+  if (failed(parser.parseOptionalRParen())) {
+    do {
+      OpAsmParser::Argument newArg;              // The "new" variable name
+      OpAsmParser::UnresolvedOperand oldOperand; // The "old" input variable
+
+      // Parse "%new"
+      if (parser.parseArgument(newArg)) {
+        return failure();
+      }
+
+      // Parse "="
+      if (parser.parseEqual()) {
+        return failure();
+      }
+
+      // Parse "%old"
+      if (parser.parseOperand(oldOperand)) {
+        return failure();
+      }
+      operands.push_back(oldOperand);
+
+      Type type = qco::QubitType::get(parser.getBuilder().getContext());
+      newArg.type = type;
+      blockArgs.push_back(newArg);
+
+    } while (succeeded(parser.parseOptionalComma()));
+
+    if (parser.parseRParen()) {
+      return failure();
+    }
+  }
+
+  // 4. Parse the Region
+  // We explicitly pass the blockArgs we just parsed so they become the entry
+  // block!
+  if (parser.parseRegion(region, blockArgs)) {
+    return failure();
+  }
+
+  return success();
+}
+
+void printTargetAliasing(OpAsmPrinter& printer, Operation* op, Region& region,
+                         OperandRange targets_in) {
+  printer << "(";
+  Block& entryBlock = region.front();
+  auto blockArgs = entryBlock.getArguments();
+
+  for (unsigned i = 0; i < targets_in.size(); ++i) {
+    if (i > 0) {
+      printer << ", ";
+    }
+    printer.printOperand(blockArgs[i]);
+    printer << " = ";
+    printer.printOperand(targets_in[i]);
+  }
+  printer << ") ";
+
+  printer.printRegion(region, /*printEntryBlockArgs=*/false);
+}
+} // end anonymous namespace
+
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//