feat: Add Program#control_flow_graph() method

pyproject.toml

@@ -26,7 +26,7 @@ lark = "^0.11.1"
 rpcq = "^3.10.0"
 networkx = ">=2.5"
 importlib-metadata = { version = ">=3.7.3,<5", python = "<3.8" }
-qcs-sdk-python = "0.17.6"
+qcs-sdk-python = "0.17.7"
 tenacity = "^8.2.2"
 types-python-dateutil = "^2.8.19"
 types-retry = "^0.9.9"

pyquil/control_flow_graph.py

@@ -0,0 +1,60 @@
+from typing import List, Optional
+from typing_extensions import Self, override
+
+from quil import program as quil_rs
+
+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 BasicBlock documentation`_ for documentation and available methods.
+
+    .. _quil BasicBlock documentation: https://rigetti.github.io/quil-rs/quil/program.html#BasicBlock
+    """
+
+    @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(quil_rs.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.
+
+    This class should not be initialized directly. Use :py:meth:~pyquil.quil.Program.control
+     flow_graph` to get a CFG for a program.
+
+    Most functionality is implemented by the `quil` package. See the `quil ControlFlowGraph documentation`_ for
+    available methods.
+
+    .. _quil ControlFlowGraph documentation: https://rigetti.github.io/quil-rs/quil/program.html#ControlFlowGraph
+    """
+
+    @classmethod
+    def _from_rs(cls, graph: quil_rs.ControlFlowGraph) -> Self:
+        return super().__new__(cls, graph)
+
+    @override
+    def basic_blocks(self) -> List[BasicBlock]:  # type: ignore[override]
+        """
+        Return a list of all the basic blocks in the control flow graph, in order of definition.
+        """
+        return [BasicBlock._from_rs(block) for block in super().basic_blocks()]

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 :py:class:`~pyquil.control_flow_graph.ControlFlowGraph` 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,59 @@
+from pyquil.control_flow_graph import BasicBlock
+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 = program.control_flow_graph()
+    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(), (type(None), AbstractInstruction))
+        assert all([isinstance(instruction, AbstractInstruction) for instruction in block.instructions()])
+        assert isinstance(block.gate_depth(1), int)