Source code for qarp.blocks._composite_block

"""User-facing CompositeBlock with measurement / target_qubit handling.

Subclass of ``qarp.blocks.CompositeBlockBase`` (which is
``qarpx.CompositeBlock`` + Python state).  Adds what the C++ class doesn't
carry: a children list wired at build, and qubit-count resolution from the
children.  Children are wired as given — controlisation is explicit through
``ControlledBlock`` (§13).
"""

from typing import List, Optional, Sequence

import qarpx as qx

from ._block import CompositeBlockBase




class CompositeBlock(CompositeBlockBase):
    """Compose a sequence of sub-blocks into a single circuit.

    The constructor's ``blocks`` are wired at ``build()``; ``add_child``
    appends one to wire at the next build.  A block that computes its own
    children in ``build_vanilla`` extends ``CompositeBlockBase`` instead.
    """

    def __init__(
        self,
        # Any qarp/qarpx block is a valid child (Simple, Composite, Controlled,
        # ...) — the runtime check below is `isinstance(block, qx.Block)`.
        blocks: Sequence[qx.Block],
        n_qubits: Optional[int] = None,
        target_qubits: Optional[List[int]] = None,
        *,
        name: str = "CompositeBlock",
    ) -> None:
        # If n_qubits not given, infer from children's target_qubits / n_qubits.
        if n_qubits is None:
            n_qubits = self._infer_n_qubits(blocks)
        super().__init__(n_qubits=n_qubits, target_qubits=target_qubits, name=name)
        self.blocks = list(blocks)
        # How many of ``self.blocks`` are already wired into the C++ children;
        # a rebuild after ``add_child`` wires only the tail (P1.9).
        self._n_wired = 0
        self._validate_inputs()

    # ── Validation ────────────────────────────────────────────────────

    def _validate_inputs(self) -> None:
        if not self.blocks and self.n_qubits is None:
            raise ValueError("CompositeBlock must have at least one block or n_qubits specified.")
        for block in self.blocks:
            # Accept any qarp / qarpx block (Simple, Composite, Controlled,
            # Measure, Reset, Conditional).  The Python-level ``Block`` alias
            # is ``SimpleBlock``, so we check the C++ base via ``qx.Block``
            # which every qarp block inherits from.
            if not isinstance(block, qx.Block):
                raise TypeError(f"All elements must be Block instances, got {type(block)}")

    @staticmethod
    def _infer_n_qubits(blocks: Sequence[qx.Block]) -> int:
        """Look at each block's target_qubits / n_qubits to size the parent.

        A ``ControlledBlock`` child already reports its full width (controls
        included), so no child is re-interpreted here.
        """
        required = 0
        for b in blocks:
            tq = getattr(b, "target_qubits", None)
            if tq is not None:
                required = max(required, max(tq) + 1 if tq else 0)
            else:
                required = max(required, getattr(b, "n_qubits", 0))
        return max(required, 1)

    # ── Build: walk children, populate self via add_child ─────────────



    def add_child(self, child: qx.Block) -> None:
        """Append a child and schedule it for wiring at the next ``build()``.

        Idempotent under rebuild: a composite built once, then given another
        child, wires only that child when built again.  Wiring is deferred —
        this method clears the Python built flag, so ``flatten()`` raises
        until the next ``build()`` re-enters ``build_vanilla``.
        """
        if not isinstance(child, qx.Block):
            raise TypeError(f"All elements must be Block instances, got {type(child)}")
        self.blocks.append(child)
        self.symbols = None  # re-published from the wider stream at build()
        self._built = False




    def build_vanilla(self) -> None:
        """Wire the not-yet-wired children (all of them on the first build).

        ``add_wired_child`` auto-materialises any pending Python-level lazy
        ops on the child, so a child created via ``set_symbols`` / ``dagger``
        is folded into a concrete ``SimpleBlock`` at composition time.
        """
        for block in self.blocks[self._n_wired :]:
            self.add_wired_child(block)
        self._n_wired = len(self.blocks)




    def unbuilt_children(self) -> Optional[List[qx.Block]]:
        """The children ``build()`` will wire, in order, readable before
        the block is built; None from a subclass that builds differently,
        whose stream cannot be read from its list."""
        if type(self).build_vanilla is not CompositeBlock.build_vanilla:
            return None
        return list(self.blocks)


    def _post_build_check(self) -> None:
        # A subclass's own build_vanilla never wires the queue add_child fills,
        # which would leave the block built and empty.
        unwired = len(self.blocks) - self._n_wired
        if unwired:
            raise TypeError(
                f"{self.name}: {unwired} child block(s) were added but never wired.  "
                "CompositeBlock.add_child only queues a child for "
                "CompositeBlock.build_vanilla; a block that computes its own "
                "children in build_vanilla extends CompositeBlockBase, whose "
                "add_child wires at once."
            )
        super()._post_build_check()