orchestrata_ultimate.refactor.codemod._csr_index

Reference › orchestrata_ultimate › refactor › codemod › _csr_index

_CSRIndex stores a directed graph as compressed sparse row arrays so that the neighbours of a node can be read quickly. It imports nothing from the project and is used by _cig and _cig_builder_ast. It is built from a node count and a list of (source, target) pairs, with a forward and a reverse array. dependencies and dependents return the successors and predecessors of a node, out_degree and in_degree count them, and transitive_closure returns all nodes reachable forward or backward with an optional depth limit.

At a glance

Class diagram

Dependency diagram

API

Classes

_CSRIndex

Line 18

Compressed Sparse Row adjacency index for directed graphs.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/refactor/codemod/_csr_index.py

"""Compressed Sparse Row (CSR) adjacency arrays for O(1) graph queries.

Stores a directed graph as two CSR matrices (forward and transpose) for
O(1) neighbor lookups in both directions. Memory layout:

    _offsets[i]   : start index in _neighbors for node i
    _offsets[i+1] : end index (exclusive)
    _neighbors[_offsets[i]:_offsets[i+1]] : neighbor node_ids
"""

from __future__ import annotations

from array import array

import codemod_fast


class _CSRIndex:
    """Compressed Sparse Row adjacency index for directed graphs."""

    __slots__ = (
        "_node_count",
        "_offsets",
        "_neighbors",
        "_transpose_offsets",
        "_transpose_neighbors",
    )

    def __init__(self, node_count: int, edges: list[tuple[int, int]]) -> None:
        self._node_count = node_count
        self._offsets, self._neighbors = self._build_csr(node_count, edges)
        self._transpose_offsets, self._transpose_neighbors = self._build_csr(
            node_count, [(to, frm) for frm, to in edges]
        )

    @staticmethod
    def _build_csr(
        node_count: int, edges: list[tuple[int, int]]
    ) -> tuple[array[int], array[int]]:
        """Build CSR arrays from edge list.

        Delegates to codemod_fast (C extension; pure-Python fallback inside
        the package if the native module is unavailable).
        """
        # codemod_fast is an untyped in-repo native extension (no .pyi stub).
        built: tuple[array[int], array[int]] = codemod_fast.csr_build(node_count, edges)
        return built

    @property
    def node_count(self) -> int:
        """Number of nodes in the graph."""
        return self._node_count

    def dependencies(self, node_id: int) -> list[int]:
        """Forward neighbors (what this node depends on)."""
        start = self._offsets[node_id]
        end = self._offsets[node_id + 1]
        return list(self._neighbors[start:end])

    def dependents(self, node_id: int) -> list[int]:
        """Reverse neighbors (what depends on this node)."""
        start = self._transpose_offsets[node_id]
        end = self._transpose_offsets[node_id + 1]
        return list(self._transpose_neighbors[start:end])

    def out_degree(self, node_id: int) -> int:
        """Number of outgoing edges for a node."""
        return self._offsets[node_id + 1] - self._offsets[node_id]

    def in_degree(self, node_id: int) -> int:
        """Number of incoming edges for a node."""
        return self._transpose_offsets[node_id + 1] - self._transpose_offsets[node_id]

    def transitive_closure(
        self, node_id: int, direction: str = "forward", max_depth: int = -1
    ) -> frozenset[int]:
        """BFS transitive closure with optional depth limit.

        Args:
            node_id: Starting node.
            direction: "forward" for dependencies, "reverse" for dependents.
            max_depth: Maximum BFS depth (-1 for unlimited).

        Returns:
            Set of all reachable node IDs (including start node).
        """
        if direction == "forward":
            offsets, neighbors = self._offsets, self._neighbors
        else:
            offsets, neighbors = (
                self._transpose_offsets,
                self._transpose_neighbors,
            )
        return frozenset(codemod_fast.csr_bfs(offsets, neighbors, node_id, max_depth))