utils: (qssa-builder) introduce builder for qssa

inconspiquous/utils/qssa_builder.py

@@ -0,0 +1,52 @@
+from xdsl.builder import Builder, ImplicitBuilder
+from xdsl.ir import Operation, SSAValue
+from dataclasses import dataclass
+
+from inconspiquous.dialects.qssa import DynGateOp, GateOp, MeasureOp
+from inconspiquous.dialects.qubit import AllocOp
+from inconspiquous.gates.core import GateAttr
+
+
+@dataclass
+class QubitRef:
+    qubit: SSAValue | None
+
+    def get(self) -> SSAValue:
+        if self.qubit is None:
+            raise ValueError("Consumed qubit was used.")
+        return self.qubit
+
+
+@dataclass
+class QSSABuilder(Builder):
+    """
+    Helper for building qssa circuits.
+    """
+
+    def gate(self, gate: GateAttr | SSAValue | Operation, *qubit_refs: QubitRef):
+        if isinstance(gate, GateAttr):
+            new_op = GateOp(gate, *(ref.get() for ref in qubit_refs))
+        else:
+            new_op = DynGateOp(gate, *(ref.get() for ref in qubit_refs))
+        if ImplicitBuilder.get() is None:
+            self.insert(new_op)
+        for ref, qubit in zip(qubit_refs, new_op.outs):
+            qubit.name_hint = ref.get().name_hint
+            ref.qubit = qubit
+
+    def alloc(self, *, name_hint: str | None = None) -> QubitRef:
+        new_op = AllocOp()
+        if ImplicitBuilder.get() is None:
+            self.insert(new_op)
+        qubit = new_op.outs[0]
+        qubit.name_hint = name_hint
+        return QubitRef(qubit)
+
+    def measure(self, ref: QubitRef, *, name_hint: str | None = None) -> SSAValue:
+        new_op = MeasureOp(ref.get())
+        if ImplicitBuilder.get() is None:
+            self.insert(new_op)
+        ref.qubit = None
+        out = new_op.out
+        out.name_hint = name_hint
+        return out

tests/builder/test_qssa_builder.py

@@ -0,0 +1,79 @@
+from xdsl.ir import Block
+from xdsl.rewriter import InsertPoint
+from inconspiquous.dialects import qssa, qubit, gate
+from inconspiquous.dialects.gate import HadamardGate
+from inconspiquous.utils.qssa_builder import QSSABuilder
+
+import pytest
+
+
+def test_qssa_builder_alloc():
+    block = Block()
+    builder = QSSABuilder(InsertPoint.at_start(block))
+
+    no_hint = builder.alloc()
+    assert isinstance(no_hint.get().owner, qubit.AllocOp)
+    assert no_hint.get().name_hint is None
+
+    hint = builder.alloc(name_hint="test")
+    assert isinstance(hint.get().owner, qubit.AllocOp)
+    assert hint.get().name_hint == "test"
+
+
+def test_qssa_gate():
+    block = Block()
+    builder = QSSABuilder(InsertPoint.at_start(block))
+
+    no_hint = builder.alloc()
+
+    builder.gate(HadamardGate(), no_hint)
+    assert isinstance(no_hint.get().owner, qssa.GateOp)
+    assert no_hint.get().name_hint is None
+
+    hint = builder.alloc(name_hint="test")
+    builder.gate(HadamardGate(), hint)
+    assert isinstance(hint.get().owner, qssa.GateOp)
+    assert hint.get().name_hint == "test"
+
+
+def test_qssa_dyn_gate():
+    block = Block()
+    builder = QSSABuilder(InsertPoint.at_start(block))
+
+    no_hint = builder.alloc()
+    gate_val = gate.ConstantGateOp(HadamardGate())
+
+    builder.gate(gate_val, no_hint)
+    assert isinstance(no_hint.get().owner, qssa.DynGateOp)
+    assert no_hint.get().name_hint is None
+
+    hint = builder.alloc(name_hint="test")
+    builder.gate(gate_val, hint)
+    assert isinstance(hint.get().owner, qssa.DynGateOp)
+    assert hint.get().name_hint == "test"
+
+
+def test_qssa_measure():
+    block = Block()
+    builder = QSSABuilder(InsertPoint.at_start(block))
+
+    q = builder.alloc()
+    c = builder.measure(q)
+    assert isinstance(c.owner, qssa.MeasureOp)
+    assert c.name_hint is None
+
+    q1 = builder.alloc()
+    c1 = builder.measure(q1, name_hint="test")
+    assert isinstance(c1.owner, qssa.MeasureOp)
+    assert c1.name_hint == "test"
+
+
+def test_double_measure():
+    block = Block()
+    builder = QSSABuilder(InsertPoint.at_start(block))
+
+    q = builder.alloc()
+    builder.measure(q)
+
+    with pytest.raises(ValueError):
+        builder.measure(q)