q3alique-codeflow/codeflow/analysis/clusterer.py

from __future__ import annotations
from codeflow.models.graph import CodeGraph
from codeflow.models.node import Node, NodeType, new_cluster_id
from codeflow.models.edge import Edge, EdgeType
from codeflow.models.flow import Flow
from collections import defaultdict
import os


def _get_env_threshold() -> int:
    try:
        return int(os.environ.get("CODEFLOW_CLUSTER_THRESHOLD", "5"))
    except (ValueError, TypeError):
        return 5


def _cluster_flow(graph: CodeGraph, flow: Flow, threshold: int) -> Flow:
    parent_map: dict[str, str] = {}
    for u, v, edge in graph.edges():
        if edge.edge_type == EdgeType.CONTAINS:
            parent_map[v] = u

    children_by_parent: dict[str, dict[NodeType, list[str]]] = defaultdict(lambda: defaultdict(list))
    for nid in flow.node_ids:
        node = graph.get_node(nid)
        if node is None:
            continue
        if node.node_type in (NodeType.ENTRY_POINT, NodeType.MODULE):
            continue
        parent = parent_map.get(nid)
        if parent is None:
            continue
        children_by_parent[parent][node.node_type].append(nid)

    new_nodes: list[Node] = []
    new_edges: list[Edge] = []
    removed: set[str] = set()

    for parent_id, type_groups in children_by_parent.items():
        for node_type, siblings in type_groups.items():
            if len(siblings) < threshold:
                continue

            cluster_id = new_cluster_id()
            type_label = node_type.value.lower()
            cluster_node = Node(
                id=cluster_id,
                node_type=NodeType.CLUSTER,
                label=f"CLUSTER: {len(siblings)} {type_label} nodes",
                file_path="",
                metadata={"members": list(siblings), "member_type": node_type.value},
            )
            new_nodes.append(cluster_node)
            new_edges.append(Edge(
                source_id=parent_id,
                target_id=cluster_id,
                edge_type=EdgeType.CONTAINS,
            ))

            for child_id in siblings:
                child = graph.get_node(child_id)
                if child is None:
                    continue
                new_edges.append(Edge(
                    source_id=cluster_id,
                    target_id=child_id,
                    edge_type=EdgeType.CLUSTER_MEMBER,
                ))

                for pred in graph.predecessors(child_id):
                    if pred not in siblings and pred != parent_id:
                        edge = graph.get_edge(pred, child_id)
                        if edge:
                            new_edges.append(Edge(
                                source_id=pred,
                                target_id=cluster_id,
                                edge_type=edge.edge_type,
                                metadata=edge.metadata.copy(),
                            ))

                for succ in graph.successors(child_id):
                    if succ not in siblings:
                        edge = graph.get_edge(child_id, succ)
                        if edge:
                            new_edges.append(Edge(
                                source_id=cluster_id,
                                target_id=succ,
                                edge_type=edge.edge_type,
                                metadata=edge.metadata.copy(),
                            ))

                removed.add(child_id)

    for node in new_nodes:
        if not graph.has_node(node.id):
            graph.add_node(node)
    for edge in new_edges:
        if graph.has_node(edge.source_id) and graph.has_node(edge.target_id):
            graph.add_edge(edge)

    for nid in list(flow.node_ids):
        if nid in removed:
            flow.node_ids.remove(nid)
        if nid in flow.source_node_ids and nid in removed:
            flow.source_node_ids.remove(nid)
        if nid in flow.sink_node_ids and nid in removed:
            flow.sink_node_ids.remove(nid)
        if nid in flow.tainted_node_ids and nid in removed:
            flow.tainted_node_ids.remove(nid)

    for node in new_nodes:
        if node.id not in flow.node_ids:
            flow.node_ids.append(node.id)

    return flow


def cluster_flows(graph: CodeGraph, flows: list[Flow], threshold: int | None = None) -> CodeGraph:
    if threshold is None:
        threshold = _get_env_threshold()
    if threshold <= 0:
        return graph

    for i, flow in enumerate(flows):
        flows[i] = _cluster_flow(graph, flow, threshold)

    return graph