apisec-code-bolt 0.1.0__py3-none-any.whl

apisec_code_bolt/analysis/flow_analysis.py

@@ -0,0 +1,1406 @@
+"""
+Flow-sensitive analysis for improved call resolution.
+
+This module implements reaching definitions analysis to provide flow-sensitive
+type bindings. Instead of accumulating ALL types a variable ever has, we track
+which definitions are "live" at each program point.
+
+ALGORITHM: Reaching Definitions with Control Flow
+1. Build a control flow graph (CFG) for each function with proper branches
+2. For each assignment, create a "definition"
+3. Compute which definitions reach each program point
+4. At each use, only consider types from reaching definitions
+
+HANDLES:
+- Sequential reassignment: x = A(); x = B(); x.method() uses only B
+- Conditional branches: if cond: x = A() else: x = B() - both types reach
+- Loops: all loop body definitions reach after loop
+- Try/except: exception handler definitions join with try block
+- Return statement analysis: tracks what functions return
+- Attribute assignments: self.attr = value
+- Cross-scope references: module-level variables used in functions
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from enum import Enum, auto
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from ..parsing.base import ParsedAssignment, ParsedFile, ParsedFunction
+
+
+logger = logging.getLogger(__name__)
+
+
+# =============================================================================
+# Control Flow Graph (CFG) Representation
+# =============================================================================
+
+
+class CFGNodeType(Enum):
+    """Types of nodes in the control flow graph."""
+
+    ENTRY = auto()  # Function entry point
+    EXIT = auto()  # Function exit point
+    ASSIGNMENT = auto()  # Variable assignment
+    ATTR_ASSIGNMENT = auto()  # Attribute assignment (self.x = ...)
+    CALL = auto()  # Function call (for tracking)
+    BRANCH = auto()  # if/while condition
+    JOIN = auto()  # Merge point after branches
+    LOOP_HEADER = auto()  # Loop entry
+    LOOP_EXIT = auto()  # Loop exit
+    RETURN = auto()  # Return statement
+    RAISE = auto()  # Raise statement
+    TRY_START = auto()  # Start of try block
+    EXCEPT_START = auto()  # Start of except block
+    FINALLY_START = auto()  # Start of finally block
+
+
+@dataclass
+class CFGNode:
+    """A node in the control flow graph."""
+
+    id: int
+    node_type: CFGNodeType
+    line: int
+
+    # For assignment nodes
+    variable: str | None = None
+    assigned_type: str | None = None
+    assigned_from_call: str | None = None
+    assigned_from_var: str | None = None
+
+    # For attribute assignments (self.attr)
+    receiver: str | None = None
+    attribute: str | None = None
+
+    # For return nodes
+    return_expression: str | None = None
+    returns_call_result: bool = False
+    returned_call: str | None = None
+
+    # Control flow
+    successors: list[int] = field(default_factory=list)
+    predecessors: list[int] = field(default_factory=list)
+
+    # For analysis - context flags
+    is_in_loop: bool = False
+    is_in_conditional: bool = False
+    is_in_try: bool = False
+    is_in_except: bool = False
+    is_in_finally: bool = False
+
+    # Branch information
+    true_branch: int | None = None  # For BRANCH nodes
+    false_branch: int | None = None  # For BRANCH nodes
+
+
+@dataclass
+class Definition:
+    """A variable definition (assignment point)."""
+
+    id: int  # Unique ID for this definition
+    variable: str  # Variable being defined
+    line: int  # Line number
+    cfg_node: int  # CFG node ID
+
+    # What type is assigned
+    assigned_type: str | None = None
+    assigned_from_call: str | None = None
+    assigned_from_var: str | None = None
+
+    # All variables the value derives from (superset of assigned_from_var).
+    # For "x = int(y)" → ["y"], "q = f'{a} {b}'" → ["a", "b"].
+    source_variables: list[str] = field(default_factory=list)
+
+    # Is this a parameter definition?
+    is_parameter: bool = False
+
+    # Is this a self/cls binding?
+    is_self_cls: bool = False
+    enclosing_class: str | None = None
+
+    # For attribute definitions (self.attr = ...)
+    is_attribute: bool = False
+    receiver: str | None = None
+    attribute: str | None = None
+
+    # Transformation metadata (for tracking what operation produced this value)
+    is_method_call: bool = False
+    is_string_interpolation: bool = False
+
+
+@dataclass
+class ReturnInfo:
+    """Information about a return statement."""
+
+    line: int
+    return_type: str | None = None
+    returns_call_result: bool = False
+    returned_call: str | None = None
+    returns_variable: bool = False
+    returned_variable: str | None = None
+    returns_literal: bool = False
+
+
+@dataclass
+class CallSiteContext:
+    """Context information about where a call occurs."""
+
+    line: int
+    in_loop: bool = False
+    in_conditional: bool = False
+    in_try: bool = False
+    in_except: bool = False
+    in_finally: bool = False
+    in_comprehension: bool = False
+    in_lambda: bool = False
+    in_with: bool = False
+
+    # Nesting depth for each context
+    loop_depth: int = 0
+    conditional_depth: int = 0
+    try_depth: int = 0
+
+
+class ControlFlowGraph:
+    """
+    A control flow graph for a function with proper branching.
+
+    Used for reaching definitions analysis.
+    """
+
+    def __init__(self, function_qname: str):
+        self.function_qname = function_qname
+        self._nodes: dict[int, CFGNode] = {}
+        self._next_id = 0
+
+        # Entry and exit nodes
+        self.entry = self._create_node(CFGNodeType.ENTRY, 0)
+        self.exit = self._create_node(CFGNodeType.EXIT, 0)
+
+        # Definitions in this function
+        self._definitions: dict[int, Definition] = {}
+        self._next_def_id = 0
+
+        # Variable -> list of definition IDs
+        self._var_definitions: dict[str, list[int]] = {}
+
+        # Attribute -> list of definition IDs (for self.x tracking)
+        self._attr_definitions: dict[tuple[str, str], list[int]] = {}
+
+        # Return statements
+        self._returns: list[ReturnInfo] = []
+
+        # Call site contexts: line -> context
+        self._call_contexts: dict[int, CallSiteContext] = {}
+
+    def _create_node(self, node_type: CFGNodeType, line: int) -> int:
+        """Create a new CFG node."""
+        node_id = self._next_id
+        self._next_id += 1
+        self._nodes[node_id] = CFGNode(
+            id=node_id,
+            node_type=node_type,
+            line=line,
+        )
+        return node_id
+
+    def add_assignment(
+        self,
+        line: int,
+        variable: str,
+        assigned_type: str | None = None,
+        assigned_from_call: str | None = None,
+        assigned_from_var: str | None = None,
+        is_parameter: bool = False,
+        is_self_cls: bool = False,
+        enclosing_class: str | None = None,
+        context_flags: dict[str, bool] | None = None,
+        source_variables: list[str] | None = None,
+        is_method_call: bool = False,
+        is_string_interpolation: bool = False,
+    ) -> int:
+        """Add an assignment node and create a definition."""
+        node_id = self._create_node(CFGNodeType.ASSIGNMENT, line)
+        node = self._nodes[node_id]
+        node.variable = variable
+        node.assigned_type = assigned_type
+        node.assigned_from_call = assigned_from_call
+        node.assigned_from_var = assigned_from_var
+
+        if context_flags:
+            node.is_in_loop = context_flags.get("in_loop", False)
+            node.is_in_conditional = context_flags.get("in_conditional", False)
+            node.is_in_try = context_flags.get("in_try", False)
+            node.is_in_except = context_flags.get("in_except", False)
+
+        # Create definition
+        def_id = self._next_def_id
+        self._next_def_id += 1
+
+        definition = Definition(
+            id=def_id,
+            variable=variable,
+            line=line,
+            cfg_node=node_id,
+            assigned_type=assigned_type,
+            assigned_from_call=assigned_from_call,
+            assigned_from_var=assigned_from_var,
+            source_variables=source_variables or [],
+            is_parameter=is_parameter,
+            is_self_cls=is_self_cls,
+            enclosing_class=enclosing_class,
+            is_method_call=is_method_call,
+            is_string_interpolation=is_string_interpolation,
+        )
+        self._definitions[def_id] = definition
+
+        if variable not in self._var_definitions:
+            self._var_definitions[variable] = []
+        self._var_definitions[variable].append(def_id)
+
+        return node_id
+
+    def add_attribute_assignment(
+        self,
+        line: int,
+        receiver: str,
+        attribute: str,
+        assigned_type: str | None = None,
+        assigned_from_call: str | None = None,
+        context_flags: dict[str, bool] | None = None,
+    ) -> int:
+        """Add an attribute assignment node (self.x = ...)."""
+        node_id = self._create_node(CFGNodeType.ATTR_ASSIGNMENT, line)
+        node = self._nodes[node_id]
+        node.receiver = receiver
+        node.attribute = attribute
+        node.assigned_type = assigned_type
+
+        if context_flags:
+            node.is_in_loop = context_flags.get("in_loop", False)
+            node.is_in_conditional = context_flags.get("in_conditional", False)
+
+        # Create definition for the attribute
+        def_id = self._next_def_id
+        self._next_def_id += 1
+
+        definition = Definition(
+            id=def_id,
+            variable=f"{receiver}.{attribute}",
+            line=line,
+            cfg_node=node_id,
+            assigned_type=assigned_type,
+            assigned_from_call=assigned_from_call,
+            is_attribute=True,
+            receiver=receiver,
+            attribute=attribute,
+        )
+        self._definitions[def_id] = definition
+
+        key = (receiver, attribute)
+        if key not in self._attr_definitions:
+            self._attr_definitions[key] = []
+        self._attr_definitions[key].append(def_id)
+
+        return node_id
+
+    def add_branch(
+        self,
+        line: int,
+        in_loop: bool = False,
+    ) -> int:
+        """Add a branch node (if/while condition)."""
+        node_id = self._create_node(CFGNodeType.BRANCH, line)
+        node = self._nodes[node_id]
+        node.is_in_loop = in_loop
+        return node_id
+
+    def add_join(self, line: int) -> int:
+        """Add a join node (merge point after branches)."""
+        return self._create_node(CFGNodeType.JOIN, line)
+
+    def add_loop_header(self, line: int) -> int:
+        """Add a loop header node."""
+        return self._create_node(CFGNodeType.LOOP_HEADER, line)
+
+    def add_return(
+        self,
+        line: int,
+        return_type: str | None = None,
+        returns_call_result: bool = False,
+        returned_call: str | None = None,
+        returns_variable: bool = False,
+        returned_variable: str | None = None,
+    ) -> int:
+        """Add a return node."""
+        node_id = self._create_node(CFGNodeType.RETURN, line)
+        node = self._nodes[node_id]
+        node.returns_call_result = returns_call_result
+        node.returned_call = returned_call
+
+        # Track return info
+        self._returns.append(
+            ReturnInfo(
+                line=line,
+                return_type=return_type,
+                returns_call_result=returns_call_result,
+                returned_call=returned_call,
+                returns_variable=returns_variable,
+                returned_variable=returned_variable,
+            )
+        )
+
+        return node_id
+
+    def add_call_context(self, line: int, context: CallSiteContext) -> None:
+        """Record the context of a call at a specific line."""
+        self._call_contexts[line] = context
+
+    def add_edge(self, from_node: int, to_node: int) -> None:
+        """Add a control flow edge."""
+        if from_node in self._nodes and to_node in self._nodes:
+            if to_node not in self._nodes[from_node].successors:
+                self._nodes[from_node].successors.append(to_node)
+            if from_node not in self._nodes[to_node].predecessors:
+                self._nodes[to_node].predecessors.append(from_node)
+
+    def set_branch_targets(
+        self,
+        branch_node: int,
+        true_target: int,
+        false_target: int,
+    ) -> None:
+        """Set the true and false branch targets for a branch node."""
+        if branch_node in self._nodes:
+            node = self._nodes[branch_node]
+            node.true_branch = true_target
+            node.false_branch = false_target
+            self.add_edge(branch_node, true_target)
+            self.add_edge(branch_node, false_target)
+
+    def get_definitions_for_var(self, variable: str) -> list[Definition]:
+        """Get all definitions for a variable."""
+        def_ids = self._var_definitions.get(variable, [])
+        return [self._definitions[d] for d in def_ids]
+
+    def get_definitions_for_attr(
+        self,
+        receiver: str,
+        attribute: str,
+    ) -> list[Definition]:
+        """Get all definitions for an attribute (e.g., self.x)."""
+        key = (receiver, attribute)
+        def_ids = self._attr_definitions.get(key, [])
+        return [self._definitions[d] for d in def_ids]
+
+    def get_call_context(self, line: int) -> CallSiteContext | None:
+        """Get the call context for a specific line."""
+        return self._call_contexts.get(line)
+
+    def get_returns(self) -> list[ReturnInfo]:
+        """Get all return statements."""
+        return self._returns
+
+
+# =============================================================================
+# Reaching Definitions Analysis
+# =============================================================================
+
+
+class ReachingDefinitions:
+    """
+    Reaching definitions analysis for a control flow graph.
+
+    Computes which definitions "reach" each program point (are not
+    killed by intervening assignments to the same variable).
+    """
+
+    def __init__(self, cfg: ControlFlowGraph):
+        self.cfg = cfg
+
+        # IN[node] = definitions reaching node entry
+        # OUT[node] = definitions leaving node
+        self._in: dict[int, set[int]] = {}
+        self._out: dict[int, set[int]] = {}
+
+        # GEN[node] = definitions created at node
+        # KILL[node] = definitions killed at node
+        self._gen: dict[int, set[int]] = {}
+        self._kill: dict[int, set[int]] = {}
+
+        # Analysis complete?
+        self._analyzed = False
+
+    def analyze(self, max_iterations: int = 100) -> int:
+        """
+        Run the reaching definitions analysis.
+
+        Uses iterative dataflow analysis with worklist algorithm.
+
+        Returns: Number of iterations
+        """
+        if self._analyzed:
+            return 0
+
+        # Initialize GEN and KILL for each node
+        self._compute_gen_kill()
+
+        # Initialize IN and OUT
+        for node_id in self.cfg._nodes:
+            self._in[node_id] = set()
+            self._out[node_id] = set()
+
+        # Worklist algorithm with reverse postorder for efficiency
+        worklist = list(self.cfg._nodes.keys())
+        iterations = 0
+
+        while worklist and iterations < max_iterations:
+            iterations += 1
+            node_id = worklist.pop(0)
+            node = self.cfg._nodes[node_id]
+
+            # IN = union of OUT of all predecessors
+            new_in: set[int] = set()
+            for pred in node.predecessors:
+                new_in.update(self._out.get(pred, set()))
+
+            self._in[node_id] = new_in
+
+            # OUT = GEN ∪ (IN - KILL)
+            old_out = self._out[node_id].copy()
+            new_out = self._gen.get(node_id, set()).copy()
+            new_out.update(new_in - self._kill.get(node_id, set()))
+            self._out[node_id] = new_out
+
+            # If OUT changed, add successors to worklist
+            if new_out != old_out:
+                for succ in node.successors:
+                    if succ not in worklist:
+                        worklist.append(succ)
+
+        self._analyzed = True
+        return iterations
+
+    def _compute_gen_kill(self) -> None:
+        """Compute GEN and KILL sets for each node."""
+        for node_id, node in self.cfg._nodes.items():
+            self._gen[node_id] = set()
+            self._kill[node_id] = set()
+
+            is_assignment = node.node_type in (
+                CFGNodeType.ASSIGNMENT,
+                CFGNodeType.ATTR_ASSIGNMENT,
+            )
+
+            if is_assignment:
+                # Get the variable being defined
+                if node.node_type == CFGNodeType.ASSIGNMENT:
+                    var_name = node.variable
+                else:
+                    var_name = f"{node.receiver}.{node.attribute}"
+
+                if var_name:
+                    # This node generates a definition
+                    for def_id, defn in self.cfg._definitions.items():
+                        if defn.cfg_node == node_id:
+                            self._gen[node_id].add(def_id)
+
+                    # This node kills other definitions of the same variable
+                    for def_id, defn in self.cfg._definitions.items():
+                        if defn.variable == var_name and defn.cfg_node != node_id:
+                            self._kill[node_id].add(def_id)
+
+    def get_reaching_definitions(
+        self,
+        variable: str,
+        at_line: int,
+    ) -> list[Definition]:
+        """
+        Get definitions of a variable that reach a given line.
+
+        This finds the nearest CFG node at or before the given line
+        and returns the definitions in its OUT set (what's live after that node).
+        """
+        if not self._analyzed:
+            self.analyze()
+
+        # Find the CFG node closest to and at/before the given line
+        best_node = None
+        best_line = -1
+
+        for node_id, node in self.cfg._nodes.items():
+            # Skip ENTRY and EXIT nodes for line-based lookup
+            if node.node_type in (CFGNodeType.ENTRY, CFGNodeType.EXIT):
+                continue
+            if node.line <= at_line and node.line > best_line:
+                best_node = node_id
+                best_line = node.line
+
+        # If no assignment node found before this line, use entry's successors
+        if best_node is None:
+            entry_successors = self.cfg._nodes[self.cfg.entry].successors
+            if entry_successors:
+                best_node = entry_successors[0]
+            else:
+                # Fall back to all definitions
+                return self.cfg.get_definitions_for_var(variable)
+
+        # Get reaching definitions AFTER this node (OUT set)
+        reaching = self._out.get(best_node, set())
+
+        # Filter to requested variable
+        result = []
+        for def_id in reaching:
+            if def_id in self.cfg._definitions:
+                defn = self.cfg._definitions[def_id]
+                if defn.variable == variable:
+                    result.append(defn)
+
+        return result
+
+    def get_reaching_attr_definitions(
+        self,
+        receiver: str,
+        attribute: str,
+        at_line: int,
+    ) -> list[Definition]:
+        """Get definitions of an attribute that reach a given line."""
+        return self.get_reaching_definitions(f"{receiver}.{attribute}", at_line)
+
+    def get_types_at_point(
+        self,
+        variable: str,
+        at_line: int,
+    ) -> set[str]:
+        """Get possible types for a variable at a specific line."""
+        definitions = self.get_reaching_definitions(variable, at_line)
+
+        types: set[str] = set()
+        for defn in definitions:
+            if defn.assigned_type:
+                types.add(defn.assigned_type)
+            if defn.is_self_cls and defn.enclosing_class:
+                types.add(defn.enclosing_class)
+
+        return types
+
+    def get_variables_deriving_from(
+        self,
+        source_var: str,
+        at_line: int,
+    ) -> set[str]:
+        """
+        Get all variables whose value may derive from source_var at this line.
+
+        Handles variable renaming: if we have uid = user_id, then uid derives
+        from user_id. Returns the transitive closure (uid, user_id, and any
+        variable assigned from those).
+
+        Uses source_variables (from parser) which captures derivation through
+        function calls: "x = int(y)" means x derives from y.
+
+        Used for taint propagation: when checking if taint passes to foo(x),
+        we check if x derives from the tainted variable.
+        """
+        if not self._analyzed:
+            self.analyze()
+
+        # Build value-flow graph: edge U -> V when V's value derives from U.
+        # Use source_variables (complete list) with assigned_from_var as fallback.
+        flow_edges: dict[str, set[str]] = {}
+        for defn in self.cfg._definitions.values():
+            sources: list[str] = []
+            if defn.source_variables:
+                sources = defn.source_variables
+            elif defn.assigned_from_var:
+                sources = [defn.assigned_from_var]
+            for src in sources:
+                if src != defn.variable:
+                    if defn.variable not in flow_edges:
+                        flow_edges[defn.variable] = set()
+                    flow_edges[defn.variable].add(src)
+
+        # Reverse graph: edge U -> {V} where V gets its value from U.
+        rev_edges: dict[str, set[str]] = {}
+        for v, sources in flow_edges.items():
+            for u in sources:
+                if u not in rev_edges:
+                    rev_edges[u] = set()
+                rev_edges[u].add(v)
+
+        # Forward reachability from source_var in the reverse graph.
+        result: set[str] = {source_var}
+        worklist = [source_var]
+        while worklist:
+            u = worklist.pop()
+            for v in rev_edges.get(u, set()):
+                if v not in result:
+                    result.add(v)
+                    worklist.append(v)
+
+        return result
+
+    def variable_derives_from(
+        self,
+        variable: str,
+        source_var: str,
+        at_line: int,
+    ) -> bool:
+        """
+        Check if variable's value may derive from source_var at this line.
+
+        Considers reaching definitions: at at_line, which defs of variable
+        reach? If any of those have source_variables or assigned_from_var
+        in the derived set, True.
+        """
+        if variable == source_var:
+            return True
+
+        derived = self.get_variables_deriving_from(source_var, at_line)
+        if variable in derived:
+            return True
+
+        reaching = self.get_reaching_definitions(variable, at_line)
+        for defn in reaching:
+            # Check source_variables (complete list) first, then
+            # fall back to assigned_from_var.
+            all_sources = defn.source_variables or (
+                [defn.assigned_from_var] if defn.assigned_from_var else []
+            )
+            for src in all_sources:
+                if src in derived or src == source_var:
+                    return True
+
+        return False
+
+    def get_derivation_chain(
+        self,
+        target_var: str,
+        source_var: str,
+        at_line: int,
+    ) -> list[Definition]:
+        """
+        Return the ordered chain of definitions that connects source_var to
+        target_var through assignments.
+
+        E.g. for source=user_id → a=int(user_id) → b=html.escape(a):
+            returns [def(a, from_call=int, from_var=user_id),
+                     def(b, from_call=html.escape, from_var=a)]
+
+        The chain is ordered from source to target. Each Definition carries
+        ``assigned_from_call`` (the transformation function) and
+        ``source_variables`` / ``assigned_from_var`` (the inputs).
+
+        Returns an empty list if no chain exists or if target == source.
+        """
+        if target_var == source_var:
+            return []
+        if not self._analyzed:
+            self.analyze()
+
+        # Build adjacency: variable → list[Definition] that defines it
+        var_defs: dict[str, list[Definition]] = {}
+        for defn in self.cfg._definitions.values():
+            var_defs.setdefault(defn.variable, []).append(defn)
+
+        # BFS backwards from target_var to source_var
+        # Each node is a variable name; edges go from defn.variable to each
+        # of its source variables (assigned_from_var / source_variables).
+        visited: set[str] = set()
+        parent: dict[str, Definition | None] = {target_var: None}
+        queue = [target_var]
+        found = False
+
+        while queue and not found:
+            current = queue.pop(0)
+            if current in visited:
+                continue
+            visited.add(current)
+
+            for defn in var_defs.get(current, []):
+                sources = defn.source_variables or (
+                    [defn.assigned_from_var] if defn.assigned_from_var else []
+                )
+                for src in sources:
+                    if src in visited:
+                        continue
+                    if src not in parent:
+                        parent[src] = defn
+                    if src == source_var:
+                        found = True
+                        break
+                    queue.append(src)
+                if found:
+                    break
+
+        if not found:
+            return []
+
+        # Reconstruct path from source_var → target_var
+        chain: list[Definition] = []
+        cur = source_var
+        while cur != target_var:
+            defn = parent.get(cur)
+            if defn is None:
+                break
+            chain.append(defn)
+            cur = defn.variable
+
+        return chain
+
+    def get_attr_types_at_point(
+        self,
+        receiver: str,
+        attribute: str,
+        at_line: int,
+    ) -> set[str]:
+        """Get possible types for an attribute at a specific line."""
+        definitions = self.get_reaching_attr_definitions(receiver, attribute, at_line)
+
+        types: set[str] = set()
+        for defn in definitions:
+            if defn.assigned_type:
+                types.add(defn.assigned_type)
+
+        return types
+
+
+# =============================================================================
+# Flow-Sensitive Binding Tracker
+# =============================================================================
+
+
+class FlowSensitiveBindings:
+    """
+    Flow-sensitive variable bindings using reaching definitions.
+
+    This is an enhancement over the basic BindingTracker that provides
+    point-specific type information.
+    """
+
+    def __init__(self):
+        # Per-function CFG and analysis
+        self._cfgs: dict[str, ControlFlowGraph] = {}
+        self._analyses: dict[str, ReachingDefinitions] = {}
+
+        # Module-level bindings (not flow-sensitive within module, but tracked)
+        self._module_types: dict[tuple[Path, str], set[str]] = {}
+
+        # Module-level attribute bindings
+        self._module_attrs: dict[tuple[Path, str, str], set[str]] = {}
+
+        # Cross-scope reference: track which module vars are used in functions
+        self._module_var_refs: dict[tuple[str, str], tuple[Path, str]] = {}
+
+        # Return type inference: function -> possible return types
+        self._return_types: dict[str, set[str]] = {}
+        self._returned_callables: dict[str, set[str]] = {}
+
+    def build_cfg_for_function(
+        self,
+        function_qname: str,
+        assignments: list[ParsedAssignment],
+        parameters: list[tuple[str, str | None]],  # (name, type_annotation)
+        enclosing_class: str | None = None,
+        control_flow_info: dict[str, Any] | None = None,
+    ) -> ControlFlowGraph:
+        """Build a CFG for a function from its assignments and control flow."""
+        cfg = ControlFlowGraph(function_qname)
+        self._cfgs[function_qname] = cfg
+
+        # Add parameter definitions (at entry)
+        last_node = cfg.entry
+        for param_name, param_type in parameters:
+            is_self = param_name in ("self", "cls") and enclosing_class is not None
+            node = cfg.add_assignment(
+                line=0,  # Parameters are at entry
+                variable=param_name,
+                assigned_type=param_type if not is_self else enclosing_class,
+                is_parameter=True,
+                is_self_cls=is_self,
+                enclosing_class=enclosing_class if is_self else None,
+            )
+            cfg.add_edge(last_node, node)
+            last_node = node
+
+        # Process control flow info if provided
+        if control_flow_info:
+            last_node = self._build_cfg_with_control_flow(
+                cfg, last_node, assignments, control_flow_info
+            )
+        else:
+            # Fallback: linear CFG (less accurate but works)
+            last_node = self._build_linear_cfg(cfg, last_node, assignments)
+
+        # Connect to exit
+        cfg.add_edge(last_node, cfg.exit)
+
+        return cfg
+
+    def register_return_statements(
+        self,
+        function_qname: str,
+        return_statements: list[Any],
+    ) -> None:
+        """
+        Register parsed return statements into the CFG so function_returns_parameter works.
+        Call after build_cfg_for_function for each function.
+        """
+        if function_qname not in self._cfgs:
+            return
+        cfg = self._cfgs[function_qname]
+        for ret in return_statements:
+            if getattr(ret, "returns_variable", False) and getattr(ret, "variable_name", None):
+                cfg.add_return(
+                    line=ret.line,
+                    returns_variable=True,
+                    returned_variable=ret.variable_name,
+                )
+            elif getattr(ret, "returns_call", False) and getattr(ret, "call_name", None):
+                cfg.add_return(
+                    line=ret.line,
+                    returns_call_result=True,
+                    returned_call=ret.call_name,
+                )
+
+    def _build_linear_cfg(
+        self,
+        cfg: ControlFlowGraph,
+        last_node: int,
+        assignments: list[ParsedAssignment],
+    ) -> int:
+        """Build a linear CFG (fallback when control flow not available)."""
+        sorted_assignments = sorted(assignments, key=lambda a: a.location.line)
+
+        for assign in sorted_assignments:
+            last_node = self._process_assignment(cfg, last_node, assign)
+
+        return last_node
+
+    def _build_cfg_with_control_flow(
+        self,
+        cfg: ControlFlowGraph,
+        last_node: int,
+        assignments: list[ParsedAssignment],
+        control_flow_info: dict[str, Any],
+    ) -> int:
+        """Build CFG with proper control flow from parsed info."""
+        # Extract control flow structures
+        if_blocks = control_flow_info.get("if_blocks", [])
+        loop_blocks = control_flow_info.get("loop_blocks", [])
+        try_blocks = control_flow_info.get("try_blocks", [])
+
+        # Build line -> context mapping
+        line_contexts: dict[int, dict[str, bool]] = {}
+
+        for if_block in if_blocks:
+            start, end = if_block.get("start", 0), if_block.get("end", 0)
+            for line in range(start, end + 1):
+                if line not in line_contexts:
+                    line_contexts[line] = {}
+                line_contexts[line]["in_conditional"] = True
+
+        for loop_block in loop_blocks:
+            start, end = loop_block.get("start", 0), loop_block.get("end", 0)
+            for line in range(start, end + 1):
+                if line not in line_contexts:
+                    line_contexts[line] = {}
+                line_contexts[line]["in_loop"] = True
+
+        for try_block in try_blocks:
+            try_start = try_block.get("try_start", 0)
+            try_end = try_block.get("try_end", 0)
+            except_ranges = try_block.get("except_ranges", [])
+            finally_range = try_block.get("finally_range")
+
+            for line in range(try_start, try_end + 1):
+                if line not in line_contexts:
+                    line_contexts[line] = {}
+                line_contexts[line]["in_try"] = True
+
+            for exc_start, exc_end in except_ranges:
+                for line in range(exc_start, exc_end + 1):
+                    if line not in line_contexts:
+                        line_contexts[line] = {}
+                    line_contexts[line]["in_except"] = True
+
+            if finally_range:
+                fin_start, fin_end = finally_range
+                for line in range(fin_start, fin_end + 1):
+                    if line not in line_contexts:
+                        line_contexts[line] = {}
+                    line_contexts[line]["in_finally"] = True
+
+        # Sort assignments by line
+        sorted_assignments = sorted(assignments, key=lambda a: a.location.line)
+
+        # Group assignments by if-else branches for proper CFG
+        # For now, use linear but with context flags
+        for assign in sorted_assignments:
+            context = line_contexts.get(assign.location.line, {})
+            last_node = self._process_assignment(cfg, last_node, assign, context)
+
+        return last_node
+
+    def _process_assignment(
+        self,
+        cfg: ControlFlowGraph,
+        last_node: int,
+        assign: ParsedAssignment,
+        context_flags: dict[str, bool] | None = None,
+    ) -> int:
+        """Process a single assignment and add it to the CFG."""
+        assigned_type = None
+        assigned_from_call = None
+        assigned_from_var = None
+
+        src_vars = getattr(assign, "source_variables", None) or []
+        is_method = getattr(assign, "is_method_call", False)
+        is_fstr = getattr(assign, "is_string_interpolation", False)
+
+        if assign.source_type == "call" and assign.source_call:
+            assigned_from_call = assign.source_call
+            assigned_type = assign.source_call
+            # Also record the primary source variable for value-flow.
+            # "x = int(y)" → assigned_from_var = "y" so the derivation
+            # graph knows x derives from y (through int()).
+            if src_vars:
+                assigned_from_var = src_vars[0]
+        elif assign.source_type == "variable" and assign.source_value:
+            assigned_from_var = assign.source_value
+        elif assign.source_type == "expression":
+            # f-strings, binary ops, etc. — track source variables for
+            # value-flow even though there is no single "source" variable.
+            if src_vars:
+                assigned_from_var = src_vars[0]
+            if assign.inferred_type:
+                assigned_type = assign.inferred_type
+        elif assign.inferred_type:
+            assigned_type = assign.inferred_type
+
+        # Check if this is an attribute assignment (self.x = ...)
+        if "." in assign.target:
+            parts = assign.target.split(".", 1)
+            if len(parts) == 2:
+                receiver, attribute = parts
+                node = cfg.add_attribute_assignment(
+                    line=assign.location.line,
+                    receiver=receiver,
+                    attribute=attribute,
+                    assigned_type=assigned_type,
+                    assigned_from_call=assigned_from_call,
+                    context_flags=context_flags,
+                )
+                cfg.add_edge(last_node, node)
+                return node
+
+        # Regular variable assignment
+        node = cfg.add_assignment(
+            line=assign.location.line,
+            variable=assign.target,
+            assigned_type=assigned_type,
+            assigned_from_call=assigned_from_call,
+            assigned_from_var=assigned_from_var,
+            context_flags=context_flags,
+            source_variables=src_vars,
+            is_method_call=is_method,
+            is_string_interpolation=is_fstr,
+        )
+        cfg.add_edge(last_node, node)
+        return node
+
+    def add_return_info(
+        self,
+        function_qname: str,
+        return_type: str | None = None,
+        returns_callable: str | None = None,
+    ) -> None:
+        """Add return type information for a function."""
+        if function_qname not in self._return_types:
+            self._return_types[function_qname] = set()
+        if return_type:
+            self._return_types[function_qname].add(return_type)
+
+        if returns_callable:
+            if function_qname not in self._returned_callables:
+                self._returned_callables[function_qname] = set()
+            self._returned_callables[function_qname].add(returns_callable)
+
+    def analyze_function(self, function_qname: str) -> None:
+        """Run reaching definitions analysis for a function."""
+        if function_qname not in self._cfgs:
+            return
+
+        cfg = self._cfgs[function_qname]
+        analysis = ReachingDefinitions(cfg)
+        analysis.analyze()
+        self._analyses[function_qname] = analysis
+
+    def analyze_all(self) -> None:
+        """Run reaching definitions for all functions."""
+        for func_qname in self._cfgs:
+            self.analyze_function(func_qname)
+
+    def get_types_at_point(
+        self,
+        variable: str,
+        function_qname: str,
+        at_line: int,
+    ) -> set[str]:
+        """
+        Get flow-sensitive types for a variable at a specific line.
+
+        This returns only the types from definitions that reach this point,
+        not all types the variable ever has.
+        """
+        if function_qname not in self._analyses:
+            # Fall back to all definitions
+            if function_qname in self._cfgs:
+                cfg = self._cfgs[function_qname]
+                return {
+                    d.assigned_type
+                    for d in cfg.get_definitions_for_var(variable)
+                    if d.assigned_type
+                }
+            return set()
+
+        analysis = self._analyses[function_qname]
+        return analysis.get_types_at_point(variable, at_line)
+
+    def get_attr_types_at_point(
+        self,
+        receiver: str,
+        attribute: str,
+        function_qname: str,
+        at_line: int,
+    ) -> set[str]:
+        """Get flow-sensitive types for an attribute at a specific line."""
+        if function_qname not in self._analyses:
+            if function_qname in self._cfgs:
+                cfg = self._cfgs[function_qname]
+                return {
+                    d.assigned_type
+                    for d in cfg.get_definitions_for_attr(receiver, attribute)
+                    if d.assigned_type
+                }
+            return set()
+
+        analysis = self._analyses[function_qname]
+        return analysis.get_attr_types_at_point(receiver, attribute, at_line)
+
+    def get_call_context(
+        self,
+        function_qname: str,
+        at_line: int,
+    ) -> CallSiteContext | None:
+        """Get the context for a call at a specific line."""
+        if function_qname in self._cfgs:
+            return self._cfgs[function_qname].get_call_context(at_line)
+        return None
+
+    def add_module_binding(
+        self,
+        file_path: Path,
+        variable: str,
+        assigned_type: str,
+    ) -> None:
+        """Add a module-level binding."""
+        key = (file_path, variable)
+        if key not in self._module_types:
+            self._module_types[key] = set()
+        self._module_types[key].add(assigned_type)
+
+    def add_module_attr_binding(
+        self,
+        file_path: Path,
+        receiver: str,
+        attribute: str,
+        assigned_type: str,
+    ) -> None:
+        """Add a module-level attribute binding."""
+        key = (file_path, receiver, attribute)
+        if key not in self._module_attrs:
+            self._module_attrs[key] = set()
+        self._module_attrs[key].add(assigned_type)
+
+    def add_cross_scope_ref(
+        self,
+        function_qname: str,
+        local_var: str,
+        file_path: Path,
+        module_var: str,
+    ) -> None:
+        """Record that a local variable references a module-level variable."""
+        self._module_var_refs[(function_qname, local_var)] = (file_path, module_var)
+
+    def get_module_types(
+        self,
+        file_path: Path,
+        variable: str,
+    ) -> set[str]:
+        """Get types for a module-level variable."""
+        return self._module_types.get((file_path, variable), set())
+
+    def get_module_attr_types(
+        self,
+        file_path: Path,
+        receiver: str,
+        attribute: str,
+    ) -> set[str]:
+        """Get types for a module-level attribute."""
+        return self._module_attrs.get((file_path, receiver, attribute), set())
+
+    def get_return_types(self, function_qname: str) -> set[str]:
+        """Get inferred return types for a function."""
+        return self._return_types.get(function_qname, set())
+
+    def get_returned_callables(self, function_qname: str) -> set[str]:
+        """Get callables returned by a function."""
+        return self._returned_callables.get(function_qname, set())
+
+    def function_returns_parameter(
+        self,
+        function_qname: str,
+        param_name: str,
+    ) -> bool:
+        """
+        Check if a function returns the given parameter (for return taint propagation).
+
+        True when:
+        - return param (exact)
+        - return param.attr (attribute access: returned_variable starts with param.)
+        - return f(param) is not yet detected (would need call-arg analysis).
+        """
+        if function_qname not in self._cfgs:
+            return False
+        cfg = self._cfgs[function_qname]
+        for ret in cfg.get_returns():
+            if not ret.returns_variable or not ret.returned_variable:
+                continue
+            if ret.returned_variable == param_name:
+                return True
+            # return param.attr or return param.x.y
+            if ret.returned_variable.startswith(param_name + "."):
+                return True
+        return False
+
+    def get_variables_deriving_from(
+        self,
+        function_qname: str,
+        source_var: str,
+        at_line: int,
+    ) -> set[str]:
+        """
+        Get variables whose value may derive from source_var at this line.
+
+        Supports variable renaming for taint propagation.
+        """
+        if function_qname not in self._analyses:
+            return {source_var}
+        return self._analyses[function_qname].get_variables_deriving_from(source_var, at_line)
+
+    def variable_derives_from(
+        self,
+        function_qname: str,
+        variable: str,
+        source_var: str,
+        at_line: int,
+    ) -> bool:
+        """Check if variable derives from source_var at this line."""
+        if function_qname not in self._analyses:
+            return variable == source_var
+        return self._analyses[function_qname].variable_derives_from(variable, source_var, at_line)
+
+    def get_derivation_chain(
+        self,
+        function_qname: str,
+        target_var: str,
+        source_var: str,
+        at_line: int,
+    ) -> list[Definition]:
+        """
+        Get the ordered chain of definitions from source_var to target_var.
+
+        Each Definition in the chain carries ``assigned_from_call`` (the
+        transformation function) and transformation metadata.
+        Returns empty list when no analysis is available or no chain exists.
+        """
+        if function_qname not in self._analyses:
+            return []
+        return self._analyses[function_qname].get_derivation_chain(target_var, source_var, at_line)
+
+    def resolve_cross_scope(
+        self,
+        function_qname: str,
+        variable: str,
+        file_path: Path,
+    ) -> set[str]:
+        """Resolve a variable that might be a cross-scope reference."""
+        key = (function_qname, variable)
+        if key in self._module_var_refs:
+            ref_file, module_var = self._module_var_refs[key]
+            return self.get_module_types(ref_file, module_var)
+
+        # Also check module level directly
+        return self.get_module_types(file_path, variable)
+
+
+# =============================================================================
+# Integration with Call Graph
+# =============================================================================
+
+
+def build_flow_sensitive_bindings(
+    parsed_files: list[ParsedFile],
+) -> FlowSensitiveBindings:
+    """
+    Build flow-sensitive bindings from parsed files.
+
+    This processes all functions and builds CFGs with reaching definitions
+    analysis for improved type resolution at specific program points.
+    """
+    bindings = FlowSensitiveBindings()
+
+    for parsed in parsed_files:
+        if not parsed.success:
+            continue
+
+        file_path = parsed.path
+
+        # Process module-level assignments
+        for assign in parsed.assignments:
+            if assign.target_qualified_name:
+                parts = assign.target_qualified_name.full.split(".")
+                # Module level: module.var (2 parts)
+                if len(parts) == 2:
+                    if assign.source_type == "call" and assign.source_call:
+                        bindings.add_module_binding(file_path, assign.target, assign.source_call)
+                    elif assign.inferred_type:
+                        bindings.add_module_binding(file_path, assign.target, assign.inferred_type)
+
+                    # Check for attribute assignment at module level
+                    if "." in assign.target:
+                        target_parts = assign.target.split(".", 1)
+                        if len(target_parts) == 2:
+                            receiver, attr = target_parts
+                            if assign.source_type == "call" and assign.source_call:
+                                bindings.add_module_attr_binding(
+                                    file_path, receiver, attr, assign.source_call
+                                )
+
+        # Process functions
+        for func in parsed.functions:
+            _process_function_for_flow(bindings, func, parsed, file_path)
+
+        # Process methods in classes
+        for cls in parsed.classes:
+            class_qname = cls.qualified_name.full
+
+            for method in cls.methods:
+                _process_function_for_flow(
+                    bindings, method, parsed, file_path, enclosing_class=class_qname
+                )
+
+    # Run analysis for all functions
+    bindings.analyze_all()
+
+    return bindings
+
+
+def _process_function_for_flow(
+    bindings: FlowSensitiveBindings,
+    func: ParsedFunction,
+    parsed: ParsedFile,
+    file_path: Path,
+    enclosing_class: str | None = None,
+) -> None:
+    """Process a function to build its CFG."""
+    qname = func.qualified_name.full
+
+    # Collect assignments in this function
+    func_assignments = []
+    for assign in parsed.assignments:
+        if assign.target_qualified_name:
+            assign_qname = assign.target_qualified_name.full
+            # Check if assignment is inside this function
+            if assign_qname.startswith(qname + ".") or assign_qname == qname:
+                func_assignments.append(assign)
+
+    # Get parameters
+    parameters = [(p.name, p.type_annotation) for p in func.parameters]
+
+    # Build control flow info from function metadata
+    control_flow_info = None
+    if hasattr(func, "control_flow_info") and func.control_flow_info:
+        control_flow_info = func.control_flow_info
+
+    # Build CFG
+    bindings.build_cfg_for_function(
+        qname,
+        func_assignments,
+        parameters,
+        enclosing_class,
+        control_flow_info,
+    )
+
+    # Register return statements so function_returns_parameter works
+    if hasattr(func, "return_statements") and func.return_statements:
+        bindings.register_return_statements(qname, func.return_statements)
+
+    # Process annotated return type
+    if func.return_type:
+        bindings.add_return_info(qname, return_type=func.return_type)
+
+    # Track return statements for type inference
+    if hasattr(func, "return_statements") and func.return_statements:
+        for ret in func.return_statements:
+            # Return calls contribute the called function's return type
+            if ret.returns_call and ret.call_name:
+                # The return type is the return type of the called function
+                # For now, we track the call name as a potential return type source
+                bindings.add_return_info(
+                    qname,
+                    return_type=f"<return:{ret.call_name}>",
+                    returns_callable=ret.call_name if ret.call_name.endswith("_factory") else None,
+                )
+
+            # Return variables - track what variable types flow to the return
+            elif ret.returns_variable and ret.variable_name:
+                # The return type comes from the variable's type at that point
+                # We mark this for later resolution when we know variable types
+                bindings.add_return_info(
+                    qname,
+                    return_type=f"<var:{ret.variable_name}>",
+                )
+
+            # Return literals have concrete types
+            elif ret.returns_literal and ret.literal_type:
+                bindings.add_return_info(qname, return_type=ret.literal_type)
+
+            # Return lambdas - track that this function returns a callable
+            elif ret.returns_lambda:
+                bindings.add_return_info(
+                    qname,
+                    returns_callable=f"{qname}.<lambda>",
+                )
+
+            # Return comprehensions have list/dict/set types
+            elif ret.returns_comprehension:
+                # Infer type from expression
+                if ret.expression_text:
+                    if ret.expression_text.startswith("{") and ":" in ret.expression_text:
+                        bindings.add_return_info(qname, return_type="dict")
+                    elif ret.expression_text.startswith("{"):
+                        bindings.add_return_info(qname, return_type="set")
+                    else:
+                        bindings.add_return_info(qname, return_type="list")
+
+            # Return None
+            elif ret.returns_none:
+                bindings.add_return_info(qname, return_type="None")