feat: Add Program#control_flow_graph() method

pyquil/control_flow_graph.py

@@ -0,0 +1,71 @@
+from typing import List, Iterable, TYPE_CHECKING, Optional
+from typing_extensions import Self, override
+
+from quil import program as quil_rs
+
+if TYPE_CHECKING:
+    from pyquil.quil import Program
+
+from pyquil.quilbase import _convert_to_py_instruction, _convert_to_py_instructions, AbstractInstruction
+
+
+class BasicBlock(quil_rs.BasicBlock):
+    """
+    Represents a basic block in the Program.
+
+    Most functionality is implemented by the `quil` package. See the
+    [quil documentation](https://rigetti.github.io/quil-rs/quil/program.html#BasicBlock) for available methods.
+    """
+
+    @classmethod
+    def _from_rs(cls, block: quil_rs.BasicBlock) -> Self:
+        return super().__new__(cls, block)
+
+    @override
+    def instructions(self) -> List[AbstractInstruction]:  # type: ignore[override]
+        return _convert_to_py_instructions(super().instructions())
+
+    @override
+    def terminator(self) -> Optional[AbstractInstruction]:  # type: ignore[override]
+        inst = super().terminator()
+        if inst is None:
+            return None
+        return _convert_to_py_instruction(super().terminator())
+
+
+class ControlFlowGraph:
+    """
+    Representation of a control flow graph (CFG) for a Quil program.
+
+    The CFG is a directed graph where each node is a basic block and each edge is a control flow transition between two
+    basic blocks.
+    """
+
+    def __init__(self, program: "Program"):
+        self._graph = program._program.control_flow_graph()
+
+    @staticmethod
+    def _from_rs(graph: quil_rs.ControlFlowGraph) -> "ControlFlowGraph":
+        from pyquil.quil import Program
+
+        py_graph = ControlFlowGraph(Program())
+        py_graph._graph = graph
+
+        return py_graph
+
+    def has_dynamic_control_flow(self) -> bool:
+        """
+        Return True if the program has dynamic control flow, i.e. contains a conditional branch instruction.
+
+        False does not imply that there is only one basic block in the program. Multiple basic blocks may have
+        non-conditional control flow among them, in which the execution order is deterministic and does not depend on
+        program state. This may be a sequence of basic blocks with fixed JUMPs or without explicit terminators.
+        """
+        return self._graph.has_dynamic_control_flow()
+
+    def basic_blocks(self) -> Iterable[BasicBlock]:
+        """
+        Return a list of all the basic blocks in the control flow graph, in order of definition.
+        """
+        for block in self._graph.basic_blocks():
+            yield BasicBlock._from_rs(block)

pyquil/quil.py

@@ -43,6 +43,7 @@
 
 from qcs_sdk.compiler.quilc import NativeQuilMetadata
 
+from pyquil.control_flow_graph import ControlFlowGraph
 from pyquil.gates import MEASURE, RESET
 from pyquil.noise import _check_kraus_ops, _create_kraus_pragmas, pauli_kraus_map
 from pyquil.quilatom import (
@@ -320,6 +321,10 @@ def inst(self, *instructions: Union[InstructionDesignator, RSProgram]) -> "Progr
 
         return self
 
+    def control_flow_graph(self) -> ControlFlowGraph:
+        """Return the control flow graph of the program."""
+        return ControlFlowGraph._from_rs(self._program.control_flow_graph())
+
     def with_loop(
         self,
         num_iterations: int,

test/unit/test_control_flow_graph.py

@@ -0,0 +1,61 @@
+from types import NoneType
+
+from pyquil.control_flow_graph import BasicBlock, ControlFlowGraph
+from pyquil.quilbase import AbstractInstruction
+from pyquil.quil import Program
+
+
+def test_control_flow_graph():
+    program = Program(
+        """
+DEFFRAME 0 "flux_tx_cz":
+    TEST: 1
+
+DEFFRAME 1 "flux_tx_iswap":
+    TEST: 1
+
+DEFFRAME 1 "flux_tx_cz":
+    TEST: 1
+
+DEFFRAME 1 "flux_tx_iswap":
+    TEST: 1
+
+DEFFRAME 2 "flux_tx_cz":
+    TEST: 1
+
+DEFFRAME 2 "flux_tx_iswap":
+    TEST: 1
+
+DEFFRAME 3 "flux_tx_cz":
+    TEST: 1
+
+DEFFRAME 3 "flux_tx_iswap":
+    TEST: 1
+
+# Simplified version
+DEFCAL CZ q0 q1:
+    FENCE q0 q1
+    SET-PHASE q0 "flux_tx_cz" 0.0
+    SET-PHASE q1 "flux_tx_iswap" 0.0
+    NONBLOCKING PULSE q0 "flux_tx_cz" erf_square(duration: 6.000000000000001e-08)
+    NONBLOCKING PULSE q1 "flux_tx_iswap" erf_square(duration: 6.000000000000001e-08)
+    SHIFT-PHASE q0 "flux_tx_cz" 1.0
+    SHIFT-PHASE q1 "flux_tx_iswap" 1.0
+    FENCE q0 q1
+
+CZ 0 1
+CZ 2 3
+CZ 0 2
+JUMP @END
+LABEL @END
+"""
+    )
+    graph = ControlFlowGraph(program)
+    assert not graph.has_dynamic_control_flow()
+    blocks = list(graph.basic_blocks())
+    assert len(blocks) == 2
+    for block in blocks:
+        assert isinstance(block, BasicBlock)
+        assert isinstance(block.terminator(), (NoneType, AbstractInstruction))
+        assert all([isinstance(instruction, AbstractInstruction) for instruction in block.instructions()])
+        assert isinstance(block.gate_depth(1), int)