orchestrata_ultimate.refactor.codemod._cig

Reference › orchestrata_ultimate › refactor › codemod › _cig

CodeIntelligenceGraph is the in-memory property graph at the heart of the codemod engine: nodes for modules, classes, functions and imports, and typed edges (CIGEdgeType, defined in _types) such as imports, calls, inheritance and tested-by. It is built by _cig_builder_ast, stores adjacency in the _CSRIndex arrays, and is queried by every module in the analysis package, by review_context and by the daemon server. It is traced to spec 110 story US-110-006 (REQ-110-022 to 024).

Lookups are by string key (resolve, node_key) or integer id. The query methods are dependencies and dependents; blast_radius, a breadth-first walk of up to a number of hops (default 3) in both directions that returns a BlastRadius of affected nodes and files; breaking_changes for a CodeChange; dependency_path between two nodes; cycles (strongly connected components); topological_order; affected_tests; coupling_score; and abstraction_level. centrality computes betweenness by the exact Brandes algorithm by default and caches the result; the docstring records that a sampled estimate was measured as inaccurate and is opt-in. pagerank_score is also cached.

At a glance

Class diagram

Dependency diagram

API

Classes

CodeIntelligenceGraph

Decorators: traces(spec='110', story='US-110-006', reqs=['REQ-110-022', 'REQ-110-023', 'REQ-110-024']) · Line 28

In-memory property graph with multihop traversal.

Source

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

"""Code Intelligence Graph with multihop traversal queries.

Provides 11 query methods over an in-memory property graph built from
the CST forest. All queries use CSR adjacency arrays for O(1) lookups.
"""

from __future__ import annotations

import itertools
from collections import deque

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.refactor.codemod._csr_index import _CSRIndex
from orchestrata_ultimate.refactor.codemod._types import (
    AbstractionLevel,
    BlastRadius,
    BreakingChange,
    CIGEdgeType,
    CIGNode,
    CodeChange,
    CouplingScore,
)


@traces(
    spec="110", story="US-110-006", reqs=["REQ-110-022", "REQ-110-023", "REQ-110-024"]
)
class CodeIntelligenceGraph:
    """In-memory property graph with multihop traversal."""

    def __init__(
        self,
        csr: _CSRIndex,
        nodes: dict[int, CIGNode],
        key_to_id: dict[str, int],
        edges_by_type: dict[CIGEdgeType, list[tuple[int, int]]],
    ) -> None:
        """Wire CodeIntelligenceGraph's composed collaborators.

        Args:
            csr: CSR adjacency index for O(1) lookups.
            nodes: Mapping of node_id to CIGNode.
            key_to_id: Mapping of qualified_name to node_id.
            edges_by_type: Edges grouped by type for filtered traversal.
        """
        self._csr = csr
        self._nodes = nodes
        self._key_to_id = key_to_id
        self._id_to_key = {v: k for k, v in key_to_id.items()}
        self._edges_by_type = edges_by_type
        self._pagerank_ranks: dict[int, float] | None = None
        #: Exact betweenness vector, computed on first `centrality()` call.
        #: Instance-scoped: the graph has no mutators, and a workspace
        #: reload builds a new CodeIntelligenceGraph, so it cannot go stale.
        self._betweenness: list[float] | None = None

    @property
    def node_count(self) -> int:
        """Total number of nodes."""
        return len(self._nodes)

    @property
    def nodes(self) -> dict[int, CIGNode]:
        """All nodes indexed by node_id."""
        return self._nodes

    @property
    def key_to_id(self) -> dict[str, int]:
        """Mapping of qualified names to node IDs."""
        return self._key_to_id

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def resolve(self, key: str) -> int:
        """Resolve a qualified name to a node ID."""
        return self._key_to_id[key]

    def node_key(self, node_id: int) -> str:
        """Resolve a node ID to its qualified name."""
        return self._id_to_key.get(node_id, f"<unknown:{node_id}>")

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def _neighbors_by_type(self, node_id: int, edge_type: CIGEdgeType) -> list[int]:
        """Get neighbors connected by a specific edge type."""
        edges = self._edges_by_type.get(edge_type, [])
        return [to for frm, to in edges if frm == node_id]

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def _reverse_neighbors_by_type(
        self, node_id: int, edge_type: CIGEdgeType
    ) -> list[int]:
        """Get nodes that connect TO this node by a specific edge type."""
        edges = self._edges_by_type.get(edge_type, [])
        return [frm for frm, to in edges if to == node_id]

    def dependencies(self, node_id: int) -> list[int]:
        """Direct imports (O(1) via CSR)."""
        return self._csr.dependencies(node_id)

    def dependents(self, node_id: int) -> list[int]:
        """Direct reverse imports (O(1) via CSR transpose)."""
        return self._csr.dependents(node_id)

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def blast_radius(self, node_id: int, hops: int = 3) -> BlastRadius:
        """BFS transitive closure over CALLS + INHERITS + IMPORTS edges.

        Args:
            node_id: Starting node.
            hops: Maximum BFS depth.

        Returns:
            BlastRadius with affected nodes, files, and impact score.
        """
        visited: set[int] = set()
        frontier = [node_id]

        for _ in range(hops):
            next_frontier: list[int] = [
                n
                for nid in frontier
                for n in itertools.chain(
                    self._csr.dependencies(nid), self._csr.dependents(nid)
                )
            ]
            next_frontier = [
                n for n in next_frontier if n not in visited and n != node_id
            ]
            visited.update(next_frontier)
            frontier = next_frontier
            if not frontier:
                break

        visited.add(node_id)
        affected_files: set[str] = set()
        for nid in visited:
            node = self._nodes.get(nid)
            if node:
                affected_files.add(node.path)

        total = len(self._nodes)
        impact = len(visited) / total if total > 0 else 0.0
        return BlastRadius(
            affected_nodes=frozenset(visited),
            affected_files=frozenset(affected_files),
            impact_score=impact,
        )

    @traces(spec="110", story="US-110-003", reqs=["REQ-110-023"])
    def breaking_changes(self, change: CodeChange) -> list[BreakingChange]:
        """Detect API-breaking changes via signature diff + downstream consumers.

        Args:
            change: The code change to analyze.

        Returns:
            List of BreakingChange records for each affected consumer.
        """
        results: list[BreakingChange] = []
        if change.entity_id is None:
            return results

        consumers = self._csr.dependents(change.entity_id)
        for consumer_id in consumers:
            consumer = self._nodes.get(consumer_id)
            if consumer is None:
                continue
            results.append(
                BreakingChange(
                    change_path=change.path,
                    consumer_path=consumer.path,
                    consumer_line=consumer.line,
                    incompatibility=change.description,
                )
            )
        return results

    def dependency_path(self, from_id: int, to_id: int) -> list[int] | None:
        """Shortest dependency path (BFS).

        Args:
            from_id: Source node.
            to_id: Target node.

        Returns:
            List of node IDs forming the path, or None if unreachable.
        """
        if from_id == to_id:
            return [from_id]

        visited: set[int] = {from_id}
        queue: deque[tuple[int, list[int]]] = deque()
        queue.append((from_id, [from_id]))

        while queue:
            current, path = queue.popleft()
            for neighbor in self._csr.dependencies(current):
                if neighbor == to_id:
                    return path + [neighbor]
                if neighbor not in visited:
                    visited.add(neighbor)
                    queue.append((neighbor, path + [neighbor]))
        return None

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def cycles(self) -> list[list[int]]:
        """All import cycles via iterative Tarjan's SCC algorithm.

        Returns:
            List of strongly connected components with size > 1.
        """
        index_counter = [0]
        stack: list[int] = []
        on_stack: set[int] = set()
        index: dict[int, int] = {}
        lowlink: dict[int, int] = {}
        result: list[list[int]] = []

        all_nodes = list(self._nodes.keys())

        def _strongconnect(v: int) -> None:
            work_stack: list[tuple[int, list[int], int]] = []
            index[v] = index_counter[0]
            lowlink[v] = index_counter[0]
            index_counter[0] += 1
            stack.append(v)
            on_stack.add(v)

            neighbors = self._csr.dependencies(v)
            work_stack.append((v, neighbors, 0))

            while work_stack:
                node, nbrs, ni = work_stack[-1]
                if ni < len(nbrs):
                    work_stack[-1] = (node, nbrs, ni + 1)
                    w = nbrs[ni]
                    if w not in index:
                        index[w] = index_counter[0]
                        lowlink[w] = index_counter[0]
                        index_counter[0] += 1
                        stack.append(w)
                        on_stack.add(w)
                        work_stack.append((w, self._csr.dependencies(w), 0))
                    elif w in on_stack:
                        lowlink[node] = min(lowlink[node], index[w])
                else:
                    if lowlink[node] == index[node]:
                        scc: list[int] = []
                        while True:
                            w = stack.pop()
                            on_stack.discard(w)
                            scc.append(w)
                            if w == node:
                                break
                        if len(scc) > 1:
                            result.append(scc)
                    work_stack.pop()
                    if work_stack:
                        parent = work_stack[-1][0]
                        lowlink[parent] = min(lowlink[parent], lowlink[node])

        for v in all_nodes:
            if v not in index:
                _strongconnect(v)

        return result

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def topological_order(self) -> list[int]:
        """Kahn's algorithm for dependency-ordered module list.

        Returns:
            Node IDs in topological order (dependencies first).
        """
        in_degree: dict[int, int] = {}
        for nid in self._nodes:
            in_degree[nid] = self._csr.in_degree(nid)

        queue: deque[int] = deque()
        for nid, deg in in_degree.items():
            if deg == 0:
                queue.append(nid)

        result: list[int] = []
        while queue:
            node = queue.popleft()
            result.append(node)
            for dep in self._csr.dependencies(node):
                in_degree[dep] -= 1
                if in_degree[dep] == 0:
                    queue.append(dep)

        return result

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def affected_tests(self, function_ids: frozenset[int]) -> frozenset[int]:
        """Tests covering changed functions via TESTED_BY reverse edges.

        Args:
            function_ids: Set of function/method node IDs that changed.

        Returns:
            Set of test node IDs that cover the changed functions.
        """
        tests: set[int] = set()
        for fid in function_ids:
            test_ids = self._reverse_neighbors_by_type(fid, CIGEdgeType.TESTED_BY)
            tests.update(test_ids)
        return frozenset(tests)

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def coupling_score(self, node_id: int) -> CouplingScore:
        """Efferent/afferent coupling + instability metric.

        Args:
            node_id: Node to analyze.

        Returns:
            CouplingScore with efferent, afferent, and instability.
        """
        efferent = len(self._csr.dependencies(node_id))
        afferent = len(self._csr.dependents(node_id))
        total = efferent + afferent
        instability = efferent / total if total > 0 else 0.0
        return CouplingScore(
            efferent=efferent,
            afferent=afferent,
            instability=instability,
        )

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def abstraction_level(self, node_id: int) -> AbstractionLevel:
        """Ratio of abstract definitions (Protocol/ABC) to concrete.

        Args:
            node_id: Module node to analyze.

        Returns:
            AbstractionLevel with abstract/concrete counts and ratio.
        """
        defined = self._neighbors_by_type(node_id, CIGEdgeType.DEFINES)
        abstract_count = 0
        concrete_count = 0

        for child_id in defined:
            child = self._nodes.get(child_id)
            if child is None or child.node_type != "Class":
                continue
            is_abstract = child.properties.get(
                "is_protocol", False
            ) or child.properties.get("is_abstract", False)
            if is_abstract:
                abstract_count += 1
            else:
                concrete_count += 1

        total = abstract_count + concrete_count
        ratio = abstract_count / total if total > 0 else 0.0
        return AbstractionLevel(
            abstract_count=abstract_count,
            concrete_count=concrete_count,
            ratio=ratio,
        )

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def centrality(
        self, top_k: int = 20, *, exact: bool = True, samples: int = 256
    ) -> list[tuple[int, float]]:
        """Betweenness centrality (Brandes algorithm, simplified).

        Defaults to EXACT Brandes over every source, computed once per
        instance and cached: measured 88.3s for the first call on the
        38,380-node repo CIG, ~0ms for every call after it. The array-based
        ``_brandes`` accumulator is what makes that affordable — the former
        per-source dict-of-dicts implementation measured 340s.

        ``exact=False`` selects pivot sampling (``samples`` evenly spaced
        deterministic source pivots, scores scaled by n/samples). It is NOT
        the default because its answers are wrong at useful sample counts:
        measured on the repo CIG against the exact top-20, strided pivots
        agree on 3/20 at samples=256 (0.63s), 12/20 at 1024 (2.6s), 16/20 at
        4096 (10.3s), 17/20 at 8192 (20.9s); random pivots do no better past
        1024. Scores are biased too — samples=256 reports a top score of
        47,150 against the exact 6,574. Sampling never reached the 90%
        top-20 agreement this function was specified to hold, and at the
        sample count that comes closest it is only 4x cheaper than being
        right. Use it only where a rough ordering is acceptable and say so
        at the call site.

        Args:
            top_k: Number of top results to return.
            exact: Exact Brandes over all sources (default). False selects
                the lossy sampled estimate described above.
            samples: Number of source pivots when ``exact`` is False.

        Returns:
            List of (node_id, centrality_score) tuples, sorted descending.
        """
        all_nodes = list(self._nodes.keys())
        n = len(all_nodes)
        if exact or n <= samples:
            if self._betweenness is None:
                self._betweenness = self._brandes(all_nodes)
            cb = self._betweenness
            scale = 1.0
        else:
            sources = [all_nodes[(i * n) // samples] for i in range(samples)]
            cb = self._brandes(sources)
            scale = n / samples

        ranked = sorted(
            ((nid, cb[nid] * scale) for nid in all_nodes),
            key=lambda x: x[1],
            reverse=True,
        )
        return ranked[:top_k]

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def _brandes(self, sources: list[int]) -> list[float]:
        """Accumulate betweenness over the given source nodes.

        Arrays indexed by node id (CSR node ids are contiguous 0..n-1),
        not per-source dict-of-dicts.

        Returns:
            Unnormalized centrality scores indexed by node id.
        """
        n = len(self._nodes)
        cb = [0.0] * n
        dependencies = self._csr.dependencies

        for s in sources:
            stack: list[int] = []
            pred: list[list[int]] = [[] for _ in range(n)]
            sigma = [0] * n
            sigma[s] = 1
            dist = [-1] * n
            dist[s] = 0
            queue: deque[int] = deque([s])

            while queue:
                v = queue.popleft()
                stack.append(v)
                next_dist = dist[v] + 1
                for w in dependencies(v):
                    if dist[w] < 0:
                        dist[w] = next_dist
                        queue.append(w)
                    if dist[w] == next_dist:
                        sigma[w] += sigma[v]
                        pred[w].append(v)

            delta = [0.0] * n
            while stack:
                w = stack.pop()
                for v in pred[w]:
                    delta[v] += (sigma[v] / sigma[w]) * (1.0 + delta[w])
                if w != s:
                    cb[w] += delta[w]

        return cb

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def pagerank_score(self, node_id: int) -> float:
        """PageRank-style importance based on incoming edges.

        The full-graph rank vector is computed lazily on the first call
        and memoized on the instance; the graph is build-time immutable
        (no mutation API), so no invalidation is needed. Named
        ``pagerank_score`` (not ``impact_score``) to avoid colliding with
        ``BlastRadius.impact_score`` on the JSON-RPC surface.

        Args:
            node_id: Node to score.

        Returns:
            PageRank score (0.0 to 1.0).
        """
        n = len(self._nodes)
        if n == 0:
            return 0.0
        if self._pagerank_ranks is None:
            self._pagerank_ranks = self._compute_pagerank()
        return self._pagerank_ranks.get(node_id, 0.0)

    @traces(spec="110", story="US-110-006", reqs=["REQ-110-023"])
    def _compute_pagerank(self) -> dict[int, float]:
        """Full-graph PageRank vector (damping 0.85, 20 iterations)."""
        n = len(self._nodes)
        damping = 0.85
        iterations = 20
        rank: dict[int, float] = {nid: 1.0 / n for nid in self._nodes}

        for _ in range(iterations):
            new_rank: dict[int, float] = {}
            for nid in self._nodes:
                incoming = self._csr.dependents(nid)
                incoming_sum = 0.0
                for src in incoming:
                    out_deg = self._csr.out_degree(src)
                    if out_deg > 0:
                        incoming_sum += rank[src] / out_deg
                new_rank[nid] = (1.0 - damping) / n + damping * incoming_sum
            rank = new_rank
        return rank