Add WireIterator for the QCO Dialect

mlir/include/mlir/Dialect/QCO/Utils/WireIterator.h

@@ -0,0 +1,169 @@
+/*
+ * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
+ * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
+ * All rights reserved.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Licensed under the MIT License
+ */
+
+#pragma once
+
+#include "mlir/Dialect/QCO/IR/QCODialect.h"
+
+#include <iterator>
+#include <llvm/ADT/TypeSwitch.h>
+#include <llvm/Support/ErrorHandling.h>
+#include <mlir/IR/Operation.h>
+#include <mlir/Support/LLVM.h>
+
+namespace mlir::qco {
+
+/**
+ * @brief A bidirectional_iterator traversing the def-use chain of a qubit wire.
+ *
+ * The iterator follows the flow of a qubit through a sequence of quantum
+ * operations while respecting the semantics of the respective operation.
+ **/
+class [[nodiscard]] WireIterator {
+public:
+  using iterator_category = std::bidirectional_iterator_tag;
+  using difference_type = std::ptrdiff_t;
+  using value_type = mlir::Operation*;
+
+  WireIterator() : op_(nullptr), qubit_(nullptr), isSentinel_(false) {}
+  explicit WireIterator(mlir::Value qubit)
+      : op_(qubit.getDefiningOp()), qubit_(qubit), isSentinel_(false) {}
+
+  /// @returns the qubit the iterator points to.
+  [[nodiscard]] mlir::Value qubit() const {
+    // A deallocation doesn't have an OpResult.
+    if (mlir::isa<qco::DeallocOp>(op_)) {
+      return nullptr;
+    }
+    return qubit_;
+  }
+
+  /// @returns the operation the iterator points to.
+  [[nodiscard]] mlir::Operation* operation() const { return op_; }
+
+  /// @returns the operation the iterator points to.
+  [[nodiscard]] mlir::Operation* operator*() const { return operation(); }
+
+  WireIterator& operator++() {
+    forward();
+    return *this;
+  }
+
+  WireIterator operator++(int) {
+    auto tmp = *this;
+    ++*this;
+    return tmp;
+  }
+
+  WireIterator& operator--() {
+    backward();
+    return *this;
+  }
+
+  WireIterator operator--(int) {
+    auto tmp = *this;
+    --*this;
+    return tmp;
+  }
+
+  bool operator==(const WireIterator& other) const {
+    return other.qubit_ == qubit_ && other.op_ == op_ &&
+           other.isSentinel_ == isSentinel_;
+  }
+
+  bool operator==([[maybe_unused]] std::default_sentinel_t s) const {
+    return isSentinel_;
+  }
+
+private:
+  /// @brief Move to the next operation on the qubit wire.
+  void forward() {
+    // If the iterator is a sentinel already, there is nothing to do.
+    if (isSentinel_) {
+      return;
+    }
+
+    // Find the user-operation of the qubit SSA value.
+    assert(qubit_.getNumUses() == 1 && "expected linear typing");
+    op_ = *(qubit_.getUsers().begin());
+
+    // A deallocation op defines the end of the qubit wire (dynamic and static).
+    if (mlir::isa<qco::DeallocOp>(op_)) {
+      isSentinel_ = true;
+      return;
+    }
+
+    if (!(mlir::isa<qco::AllocOp, qco::StaticOp>(op_))) {
+      // Find the output from the input qubit SSA value.
+      mlir::TypeSwitch<mlir::Operation*>(op_)
+          .Case<qco::UnitaryOpInterface>([&](qco::UnitaryOpInterface op) {
+            qubit_ = op.getOutputForInput(qubit_);
+          })
+          .Case<qco::MeasureOp>(
+              [&](qco::MeasureOp op) { qubit_ = op.getQubitOut(); })
+          .Case<qco::ResetOp>(
+              [&](qco::ResetOp op) { qubit_ = op.getQubitOut(); })
+          .Default([&](mlir::Operation* op) {
+            report_fatal_error("unknown op in def-use chain: " +
+                               op->getName().getStringRef());
+          });
+    }
+  }
+
+  /// @brief Move to the previous operation on the qubit wire.
+  void backward() {
+    // If the iterator is a sentinel, reactivate the iterator.
+    if (isSentinel_) {
+      isSentinel_ = false;
+      return;
+    }
+
+    // For deallocations, qubit_ is an OpOperand. Hence, only get the def-op.
+    if (mlir::isa<DeallocOp>(op_)) {
+      op_ = qubit_.getDefiningOp();
+      return;
+    }
+
+    // Allocations or static definitions define the start of the qubit wire.
+    // Consequently, stop and early exit.
+    if (mlir::isa<qco::AllocOp, StaticOp>(op_)) {
+      return;
+    }
+
+    // Find the input from the output qubit SSA value.
+    mlir::TypeSwitch<mlir::Operation*>(op_)
+        .Case<qco::UnitaryOpInterface>([&](qco::UnitaryOpInterface op) {
+          qubit_ = op.getInputForOutput(qubit_);
+        })
+        .Case<qco::MeasureOp>(
+            [&](qco::MeasureOp op) { qubit_ = op.getQubitIn(); })
+        .Case<qco::ResetOp>([&](qco::ResetOp op) { qubit_ = op.getQubitIn(); })
+        .Default([&](mlir::Operation* op) {
+          report_fatal_error("unknown op in def-use chain: " +
+                             op->getName().getStringRef());
+        });
+
+    // Get the operation that produces the qubit value.
+    // If the current qubit SSA value is a BlockArgument (no defining op), stop.
+    op_ = qubit_.getDefiningOp();
+    if (op_ == nullptr) {
+      return;
+    }
+  }
+
+  mlir::Operation* op_;
+  mlir::Value qubit_;
+  bool isSentinel_;
+};
+
+static_assert(std::bidirectional_iterator<WireIterator>);
+static_assert(std::sentinel_for<std::default_sentinel_t, WireIterator>,
+              "std::default_sentinel_t must be a sentinel for WireIterator.");
+} // namespace mlir::qco

mlir/unittests/CMakeLists.txt

@@ -13,4 +13,5 @@ add_custom_target(mqt-core-mlir-unittests)
 
 add_dependencies(
   mqt-core-mlir-unittests mqt-core-mlir-compiler-pipeline-test mqt-core-mlir-qco-dialect-test
-  mqt-core-mlir-dialect-qco-ir-modifiers-test mqt-core-mlir-dialect-utils-test)
+  mqt-core-mlir-dialect-qco-ir-modifiers-test mqt-core-mlir-qco-dialect-utils-test
+  mqt-core-mlir-dialect-utils-test)

mlir/unittests/Dialect/QCO/CMakeLists.txt

@@ -7,3 +7,4 @@
 # Licensed under the MIT License
 
 add_subdirectory(IR)
+add_subdirectory(Utils)

mlir/unittests/Dialect/QCO/Utils/CMakeLists.txt

@@ -0,0 +1,15 @@
+# Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
+# Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
+# All rights reserved.
+#
+# SPDX-License-Identifier: MIT
+#
+# Licensed under the MIT License
+
+add_executable(mqt-core-mlir-qco-dialect-utils-test test_wireiterator.cpp)
+
+target_link_libraries(
+  mqt-core-mlir-qco-dialect-utils-test PRIVATE GTest::gtest_main MLIRQCODialect MLIRLLVMDialect
+                                               MLIRQCOProgramBuilder)
+
+gtest_discover_tests(mqt-core-mlir-qco-dialect-utils-test)

mlir/unittests/Dialect/QCO/Utils/test_wireiterator.cpp

@@ -0,0 +1,132 @@
+/*
+ * Copyright (c) 2023 - 2026 Chair for Design Automation, TUM
+ * Copyright (c) 2025 - 2026 Munich Quantum Software Company GmbH
+ * All rights reserved.
+ *
+ * SPDX-License-Identifier: MIT
+ *
+ * Licensed under the MIT License
+ */
+
+#include "mlir/Dialect/QCO/Builder/QCOProgramBuilder.h"
+#include "mlir/Dialect/QCO/IR/QCODialect.h"
+#include "mlir/Dialect/QCO/Utils/WireIterator.h"
+
+#include <gtest/gtest.h>
+#include <iterator>
+#include <llvm/Support/Debug.h>
+#include <memory>
+#include <mlir/Dialect/Arith/IR/Arith.h>
+#include <mlir/Dialect/ControlFlow/IR/ControlFlow.h>
+#include <mlir/Dialect/Func/IR/FuncOps.h>
+#include <mlir/Dialect/LLVMIR/LLVMDialect.h>
+#include <mlir/Dialect/SCF/IR/SCF.h>
+#include <mlir/IR/DialectRegistry.h>
+#include <mlir/IR/MLIRContext.h>
+#include <utility>
+
+using namespace mlir;
+
+namespace {
+class WireIteratorTest : public testing::TestWithParam<bool> {
+protected:
+  void SetUp() override {
+    DialectRegistry registry;
+    registry
+        .insert<qco::QCODialect, arith::ArithDialect, cf::ControlFlowDialect,
+                func::FuncDialect, scf::SCFDialect, LLVM::LLVMDialect>();
+
+    context = std::make_unique<MLIRContext>();
+    context->appendDialectRegistry(registry);
+    context->loadAllAvailableDialects();
+  }
+
+  std::unique_ptr<MLIRContext> context;
+};
+} // namespace
+
+TEST_P(WireIteratorTest, MixedUse) {
+  const bool isDynamic = GetParam();
+
+  // Build circuit.
+  qco::QCOProgramBuilder builder(context.get());
+  builder.initialize();
+  const auto q00 = isDynamic ? builder.allocQubit() : builder.staticQubit(0);
+  const auto q10 = isDynamic ? builder.allocQubit() : builder.staticQubit(1);
+  const auto q01 = builder.h(q00);
+  const auto [q02, q11] = builder.cx(q01, q10);
+  const auto [q03, c0] = builder.measure(q02);
+  const auto q04 = builder.reset(q03);
+  builder.dealloc(q04);
+  builder.dealloc(q11);
+
+  // Setup WireIterator.
+  auto module = builder.finalize();
+  qco::WireIterator it(q00);
+
+  //
+  // Test: Forward Iteration
+  //
+
+  ASSERT_EQ(it.operation(), q00.getDefiningOp()); // qco.alloc
+  ASSERT_EQ(it.qubit(), q00);
+
+  ++it;
+  ASSERT_EQ(it.operation(), q01.getDefiningOp()); // qco.h
+  ASSERT_EQ(it.qubit(), q01);
+
+  ++it;
+  ASSERT_EQ(it.operation(), q02.getDefiningOp()); // qco.ctrl
+  ASSERT_EQ(it.qubit(), q02);
+
+  ++it;
+  ASSERT_EQ(it.operation(), q03.getDefiningOp()); // qco.measure
+  ASSERT_EQ(it.qubit(), q03);
+
+  ++it;
+  ASSERT_EQ(it.operation(), q04.getDefiningOp()); // qco.reset
+  ASSERT_EQ(it.qubit(), q04);
+
+  ++it;
+  ASSERT_EQ(it.operation(), *(q04.getUsers().begin())); // qco.dealloc
+  ASSERT_EQ(it.qubit(), nullptr);
+
+  ++it;
+  ASSERT_EQ(it, std::default_sentinel);
+
+  ++it;
+  ASSERT_EQ(it, std::default_sentinel);
+
+  //
+  // Test: Backward Iteration
+  //
+
+  --it;
+  ASSERT_EQ(it.operation(), *(q04.getUsers().begin())); // qco.dealloc
+  ASSERT_EQ(it.qubit(), nullptr);
+
+  --it;
+  ASSERT_EQ(it.operation(), q04.getDefiningOp()); // qco.reset
+  ASSERT_EQ(it.qubit(), q04);
+
+  --it;
+  ASSERT_EQ(it.operation(), q03.getDefiningOp()); // qco.measure
+  ASSERT_EQ(it.qubit(), q03);
+
+  --it;
+  ASSERT_EQ(it.operation(), q02.getDefiningOp()); // qco.ctrl
+  ASSERT_EQ(it.qubit(), q02);
+
+  --it;
+  ASSERT_EQ(it.operation(), q01.getDefiningOp()); // qco.h
+  ASSERT_EQ(it.qubit(), q01);
+
+  --it;
+  ASSERT_EQ(it.operation(), q00.getDefiningOp()); // qco.alloc or qco.static
+  ASSERT_EQ(it.qubit(), q00);
+}
+
+INSTANTIATE_TEST_SUITE_P(DynamicAndStatic, WireIteratorTest, ::testing::Bool(),
+                         [](const ::testing::TestParamInfo<bool>& info) {
+                           return info.param ? "Dynamic" : "Static";
+                         });