link_graph.py

"""
link_graph.py — CrawlGraph
============================
Thread-safe directed link graph and PageRank.

DATA STRUCTURE
  An adjacency dict: { src_url: set(dst_urls) }
  Nodes are URLs; edges are hyperlinks (src → dst).
  We use a plain dict of sets rather than pulling in networkx — it keeps
  the dependency footprint small and lets the interview conversation focus
  on our synchronization choices.

SYNCHRONIZATION
  All reads and writes go through a threading.RLock (reentrant lock).
  RLock vs Lock: if a method that holds the lock calls another method that
  also tries to acquire it, Lock would deadlock; RLock lets the same thread
  re-enter. This is defensive — current code doesn't nest lock calls, but
  it's the right choice for a shared mutable object that might be subclassed.

PAGERANK
  Computed on a point-in-time snapshot (copy of the adjacency dict) taken
  under the lock. The actual iterative computation runs without holding the
  lock, so the crawler can continue writing edges during PageRank.
"""

import json
import threading
from collections import defaultdict
from typing import Dict, List, Set


class LinkGraph:
    def __init__(self) -> None:
        # adjacency: src → set of dst URLs
        self._adj: Dict[str, Set[str]] = defaultdict(set)
        self._lock = threading.RLock()

    # ── Write operations ──────────────────────────────────────────────────────

    def add_edges(self, src: str, destinations: Set[str]) -> None:
        """
        Record that src links to every URL in destinations.
        Ensures every dst is also a node (even if it has no outgoing links yet),
        so PageRank accounts for all discovered URLs, not just crawled ones.
        """
        with self._lock:
            for dst in destinations:
                self._adj[src].add(dst)
                if dst not in self._adj:
                    self._adj[dst]   # insert empty set via defaultdict

    # ── Read operations ───────────────────────────────────────────────────────

    def node_count(self) -> int:
        with self._lock:
            return len(self._adj)

    def edge_count(self) -> int:
        with self._lock:
            return sum(len(dsts) for dsts in self._adj.values())

    def _snapshot(self) -> Dict[str, List[str]]:
        """Return a plain-dict copy. Caller must hold the lock."""
        return {src: list(dsts) for src, dsts in self._adj.items()}

    # ── PageRank ──────────────────────────────────────────────────────────────

    def pagerank(
        self,
        damping: float = 0.85,
        max_iter: int = 100,
        tol: float = 1e-6,
    ) -> Dict[str, float]:
        """
        Iterative PageRank (power iteration).

        Formula:
          PR(u) = (1 - d) / N  +  d * Σ_{v→u} PR(v) / out_degree(v)

        Dangling nodes (no outgoing links) would leak rank out of the system.
        We redistribute their total rank evenly across all nodes each iteration,
        equivalent to adding a teleportation edge from dangling nodes to every
        node.

        Convergence: stop when the L1 norm of the rank delta < tol.
        """
        # Take a consistent snapshot under the lock; compute without the lock
        with self._lock:
            adj = self._snapshot()

        nodes = list(adj.keys())
        N = len(nodes)
        if N == 0:
            return {}

        out_degree: Dict[str, int] = {n: len(adj[n]) for n in nodes}

        # Reverse adjacency: who points TO each node?
        in_links: Dict[str, List[str]] = defaultdict(list)
        for src, dsts in adj.items():
            for dst in dsts:
                in_links[dst].append(src)

        # Uniform initialization
        ranks: Dict[str, float] = {n: 1.0 / N for n in nodes}

        for _ in range(max_iter):
            # Total rank held by dangling nodes (they have nowhere to send it)
            dangling_rank = sum(
                ranks[n] for n in nodes if out_degree[n] == 0
            )

            new_ranks: Dict[str, float] = {}
            for node in nodes:
                # Rank flowing in from pages that link here
                in_flow = sum(
                    ranks[v] / out_degree[v]
                    for v in in_links.get(node, [])
                    if out_degree[v] > 0
                )
                # Dangling rank redistributed uniformly
                new_ranks[node] = (
                    (1 - damping) / N
                    + damping * (in_flow + dangling_rank / N)
                )

            # L1 convergence check
            delta = sum(abs(new_ranks[n] - ranks[n]) for n in nodes)
            ranks = new_ranks
            if delta < tol:
                break

        return ranks

    # ── Export ────────────────────────────────────────────────────────────────

    def export_json(
        self,
        path: str,
        pageranks: Dict[str, float] | None = None,
    ) -> None:
        """
        Write the graph to a JSON file compatible with D3.js force layouts.
        Format: { "nodes": [...], "edges": [...] }
        """
        with self._lock:
            adj = self._snapshot()

        data = {
            "nodes": [
                {
                    "id":       url,
                    "pagerank": round(pageranks.get(url, 0.0), 8)
                    if pageranks else None,
                }
                for url in adj
            ],
            "edges": [
                {"source": src, "target": dst}
                for src, dsts in adj.items()
                for dst in dsts
            ],
        }
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)