✨ Add replaceWithAdditionalQubits utility function to qco.if

CHANGELOG.md

@@ -44,9 +44,9 @@ releases may include breaking changes.
   [#1626], [#1627], [#1635], [#1638], [#1673], [#1675], [#1700], [#1710],
   [#1717], [#1728], [#1730], [#1749], [#1751], [#1762], [#1765], [#1780],
   [#1781], [#1782], [#1787], [#1806], [#1807], [#1808], [#1823], [#1824],
-  [#1830], [#1869]) ([**@burgholzer**], [**@denialhaag**], [**@taminob**],
-  [**@DRovara**], [**@li-mingbao**], [**@Ectras**], [**@MatthiasReumann**],
-  [**@simon1hofmann**])
+  [#1830], [#1869], [#1872]) ([**@burgholzer**], [**@denialhaag**],
+  [**@taminob**], [**@DRovara**], [**@li-mingbao**], [**@Ectras**],
+  [**@MatthiasReumann**], [**@simon1hofmann**])
 
 ### Changed
 
@@ -622,6 +622,7 @@ changelogs._
 <!-- PR links -->
 
 [#1873]: https://github.com/munich-quantum-toolkit/core/pull/1873
+[#1872]: https://github.com/munich-quantum-toolkit/core/pull/1872
 [#1869]: https://github.com/munich-quantum-toolkit/core/pull/1869
 [#1850]: https://github.com/munich-quantum-toolkit/core/pull/1850
 [#1849]: https://github.com/munich-quantum-toolkit/core/pull/1849

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

@@ -1325,6 +1325,13 @@ def IfOp
         /// Return the yielded value that corresponds to the given argument
         /// for the else-block, or `nullptr` on failure.
         OpOperand* getTiedElseYieldedValue(BlockArgument bbArg);
+
+        /// Append the specified additional "qubit" operands: replace this
+        /// if-op with a new if-op that has the additional qubit operands.
+        /// The operands can be of qubit or qtensor type.
+        /// The branch bodies of this if-op are moved over to the new if-op.
+        /// The newly added qubits are yielded from each branch.
+        IfOp replaceWithAdditionalQubits(RewriterBase& rewriter, ValueRange addons);
     }];
 
   let extraClassDefinition = [{

mlir/lib/Dialect/QCO/IR/SCF/IfOp.cpp

@@ -325,3 +325,43 @@ OpOperand* IfOp::getTiedElseYieldedValue(BlockArgument bbArg) {
   }
   return &elseYield().getTargetsMutable()[bbArg.getArgNumber()];
 }
+
+IfOp IfOp::replaceWithAdditionalQubits(RewriterBase& rewriter,
+                                       ValueRange addons) {
+  if (addons.empty()) {
+    return *this;
+  }
+
+  SmallVector<Value> allQubits;
+  allQubits.reserve(getQubits().size() + addons.size());
+  allQubits.append(getQubits().begin(), getQubits().end());
+  allQubits.append(addons.begin(), addons.end());
+  const auto allQubitTypes = ValueRange(allQubits).getTypes();
+
+  auto newIfOp = create(rewriter, getLoc(), getCondition(), allQubits);
+
+  const auto rewriteRegion = [&rewriter, &allQubitTypes,
+                              &addons](Region& oldRegion, Region& newRegion) {
+    auto* oldBlock = &oldRegion.front();
+    const auto numOldArgs = oldBlock->getNumArguments();
+    auto* newBlock = rewriter.createBlock(&newRegion, {}, allQubitTypes);
+    const auto oldArgs = newBlock->getArguments().take_front(numOldArgs);
+    const auto addonArgs = newBlock->getArguments().drop_front(numOldArgs);
+
+    rewriter.mergeBlocks(oldBlock, newBlock, oldArgs);
+
+    auto yield = cast<YieldOp>(newBlock->getTerminator());
+    SmallVector<Value> yieldedValues;
+    yieldedValues.reserve(yield.getTargets().size() + addons.size());
+    yieldedValues.append(yield.getTargets().begin(), yield.getTargets().end());
+    yieldedValues.append(addonArgs.begin(), addonArgs.end());
+    rewriter.replaceOpWithNewOp<YieldOp>(yield, yieldedValues);
+  };
+
+  rewriteRegion(getThenRegion(), newIfOp.getThenRegion());
+  rewriteRegion(getElseRegion(), newIfOp.getElseRegion());
+
+  rewriter.eraseOp(*this);
+
+  return newIfOp;
+}