pipeline

def classify_direct(component: Graph, layout: ParamLayout) -> tuple[int, bool] | None:
    """Return ``(basis index, flip)`` if the component's only output is one basis bit.

    That is the case when the component is a classical output boundary joined by
    a Hadamard edge to a classical Z spider whose phase is 0 or π plus one
    error-basis parameter.
    """
    outputs = component.outputs()
    if len(outputs) != 1 or component.num_vertices() != 2:
        return None
    (out,) = outputs
    (spider,) = component.neighbors(out)
    # NOTE: zxparam never lets a classical boundary touch a quantum vertex, so a
    # classical output implies a classical spider.
    if component.vertex_domain(out) == Domain.QUANTUM:
        return None
    if component.vertex_type(spider) != VertexType.Z:
        return None
    if component.edge_type(out, spider) != EdgeType.HADAMARD:
        return None
    phase = component.phase(spider)
    if len(phase.parameters) != 1 or phase.value not in (0, 1):
        return None
    (param,) = phase.parameters
    kind, basis_index = layout.kind(param)
    if kind != "f":
        return None
    return basis_index, phase.value == 1

def compile_program(
    prepared: SamplingGraph,
    *,
    mode: DecompositionMode,
    strategy: SamplingStrategy,
) -> CompiledProgram:
    """Compile a prepared graph into an executable sampling program.

    The graph is split into connected components. A component whose single
    output equals one error-basis bit takes the direct path. With
    ``strategy="active_statevector"`` every other component becomes an action tape for
    the statevector executor. Otherwise its outputs are plugged (all prefixes
    in ``sequential`` mode, none and all in ``joint`` mode), each plugged graph
    is decomposed into stabilizer terms, and the terms are compiled into JAX
    arrays.

    Args:
        prepared: The prepared graph from ``prepare_graph``.
        mode: ``"sequential"`` for sampling, ``"joint"`` for probabilities.
        strategy: ``"active_statevector"``, or a stabilizer-rank decomposition
            strategy.

    Returns:
        A CompiledProgram ready for sampling.

    Raises:
        ValueError: If ``strategy="active_statevector"`` is combined with ``mode="joint"``.

    """
    if strategy == "active_statevector":
        if mode == "joint":
            raise ValueError(
                "strategy='active_statevector' samples outcomes and cannot compute "
                f"probabilities. {STABILIZER_RANK_HINT}"
            )
        compile_component = compile_statevector_component
    else:
        compile_component = partial(_compile_component, mode=mode, strategy=strategy)
    graph = prepared.graph
    layout = prepared.layout

    direct_entries: list[tuple[int, int, bool]] = []
    compiled_components: list[CompiledComponent | StatevectorComponent] = []
    compiled_output_order: list[int] = []

    components = sorted(graph.connected_components(), key=lambda c: len(c.outputs()))
    for component in components:
        output_indices = [_port(component, v) for v in component.outputs()]
        direct = classify_direct(component, layout)
        if direct is not None:
            basis_index, flip = direct
            direct_entries.append((output_indices[0], basis_index, flip))
            continue
        compiled_components.append(compile_component(component, output_indices, layout))
        compiled_output_order.extend(output_indices)

    # Sort direct entries by output index so that — together with the output
    # prioritisation in transform_error_basis — the concatenated layout often
    # matches the original output order, sparing a reindex at sample time.
    direct_entries.sort()
    output_order = np.array(
        [e[0] for e in direct_entries] + compiled_output_order, dtype=np.int32
    )
    reindex = np.argsort(output_order)
    is_identity = np.array_equal(reindex, np.arange(len(output_order)))

    return CompiledProgram(
        components=tuple(compiled_components),
        direct_f_indices=jnp.array([e[1] for e in direct_entries], dtype=jnp.int32),
        direct_flips=jnp.array([e[2] for e in direct_entries], dtype=jnp.bool_),
        output_order=jnp.asarray(output_order),
        output_reindex=None if is_identity else jnp.asarray(reindex),
        num_outputs=prepared.num_outputs,
        num_detectors=prepared.num_detectors,
    )