🐛 Stop QCO wire traversal at unknown carriers

mlir/lib/Dialect/QCO/Utils/WireIterator.cpp

@@ -90,10 +90,7 @@ void WireIterator::forward() {
         .Case<IndexSwitchOp>([&](IndexSwitchOp op) {
           qubit_ = op.getTiedResult(&(*qubit_.use_begin()));
         })
-        .Default([&](Operation* op) {
-          llvm::reportFatalInternalError("unknown op in def-use chain: " +
-                                         op->getName().getStringRef());
-        });
+        .Default([&](Operation*) { isFinal_ = true; });
   }
 }
 
@@ -125,6 +122,7 @@ void WireIterator::backward() {
   }
 
   // Find the input from the output qubit SSA value.
+  bool reachedBoundary = false;
   TypeSwitch<Operation*>(op_)
       .Case<UnitaryOpInterface>(
           [&](UnitaryOpInterface op) { qubit_ = op.getInputForOutput(qubit_); })
@@ -163,10 +161,13 @@ void WireIterator::backward() {
         }
         llvm::reportFatalInternalError("expected result lookup");
       })
-      .Default([&](Operation* op) {
-        llvm::reportFatalInternalError("unknown op in def-use chain: " +
-                                       op->getName().getStringRef());
-      });
+      .Default([&](Operation*) { reachedBoundary = true; });
+
+  if (reachedBoundary) {
+    op_ = nullptr;
+    isFinal_ = false;
+    return;
+  }
 
   // Get the operation that produces the qubit value.
   // If the current qubit SSA value is a BlockArgument (no defining op), the

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

@@ -17,6 +17,7 @@
 #include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/Dialect/Func/IR/FuncOps.h>
 #include <mlir/Dialect/SCF/IR/SCF.h>
+#include <mlir/IR/BuiltinOps.h>
 #include <mlir/IR/DialectRegistry.h>
 #include <mlir/IR/MLIRContext.h>
 #include <mlir/Support/LLVM.h>
@@ -264,6 +265,35 @@ TEST_P(WireIteratorTest, FunctionReturnTerminatesTraversal) {
   ASSERT_EQ(it.qubit(), output);
 }
 
+TEST_P(WireIteratorTest, UnknownCarrierTerminatesTraversal) {
+  OpBuilder builder(context.get());
+  const auto location = builder.getUnknownLoc();
+  auto module = ModuleOp::create(location);
+  builder.setInsertionPointToStart(module.getBody());
+  auto function = func::FuncOp::create(builder, location, "main",
+                                       builder.getFunctionType({}, {}));
+  Block* body = function.addEntryBlock();
+  builder.setInsertionPointToStart(body);
+
+  auto source = qco::AllocOp::create(builder, location).getResult();
+  auto carrier = UnrealizedConversionCastOp::create(
+      builder, location, TypeRange{source.getType()}, ValueRange{source});
+  auto carried = carrier.getResult(0);
+  qco::SinkOp::create(builder, location, carried);
+  func::ReturnOp::create(builder, location);
+
+  qco::WireIterator forward(source);
+  ++forward;
+  EXPECT_EQ(forward.operation(), carrier.getOperation());
+  ++forward;
+  EXPECT_EQ(forward, std::default_sentinel);
+
+  qco::WireIterator backward(carried);
+  --backward;
+  EXPECT_EQ(backward.operation(), nullptr);
+  EXPECT_EQ(backward.qubit(), carried);
+}
+
 INSTANTIATE_TEST_SUITE_P(DynamicAndStatic, WireIteratorTest, ::testing::Bool(),
                          [](const ::testing::TestParamInfo<bool>& info) {
                            return info.param ? "Dynamic" : "Static";