memory_debug.py

"""Memory debugging utilities for tracking routing memory usage."""
from __future__ import annotations

import sys
from typing import Dict, Any, Optional

# Try to import memory tracking libraries (graceful fallback)
_HAS_RESOURCE = False
_HAS_PSUTIL = False

try:
    import resource
    _HAS_RESOURCE = True
except ImportError:
    pass

try:
    import psutil
    _HAS_PSUTIL = True
except ImportError:
    pass


def get_process_memory_mb() -> float:
    """Get current process memory usage in MB.

    Tries multiple methods in order of reliability:
    1. psutil (most accurate, cross-platform)
    2. resource (Unix only)
    """
    if _HAS_PSUTIL:
        process = psutil.Process()
        return process.memory_info().rss / (1024 * 1024)

    if _HAS_RESOURCE:
        # maxrss is in KB on Linux, bytes on macOS
        maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        if sys.platform == 'darwin':
            return maxrss / (1024 * 1024)  # bytes to MB
        return maxrss / 1024  # KB to MB

    return 0.0  # Unable to measure


def format_memory_stats(label: str, memory_mb: float, delta_mb: Optional[float] = None) -> str:
    """Format memory statistics for display."""
    if delta_mb is not None and abs(delta_mb) > 0.1:
        sign = "+" if delta_mb > 0 else ""
        return f"[MEMORY] {label}: {memory_mb:.1f} MB ({sign}{delta_mb:.1f} MB)"
    return f"[MEMORY] {label}: {memory_mb:.1f} MB"


def estimate_net_obstacles_cache_mb(cache: Dict) -> float:
    """Estimate memory size of net_obstacles_cache (Dict[int, NetObstacleData])."""
    if not cache:
        return 0.0

    total = sys.getsizeof(cache)
    for net_id, data in cache.items():
        total += sys.getsizeof(net_id)
        total += sys.getsizeof(data)
        if hasattr(data, 'blocked_cells'):
            # numpy arrays: use nbytes for actual memory usage
            if hasattr(data.blocked_cells, 'nbytes'):
                total += data.blocked_cells.nbytes
            else:
                total += sys.getsizeof(data.blocked_cells)
                total += len(data.blocked_cells) * 24  # tuple of 3 ints (legacy)
        if hasattr(data, 'blocked_vias'):
            if hasattr(data.blocked_vias, 'nbytes'):
                total += data.blocked_vias.nbytes
            else:
                total += sys.getsizeof(data.blocked_vias)
                total += len(data.blocked_vias) * 16  # tuple of 2 ints (legacy)
    return total / (1024 * 1024)


def estimate_track_proximity_cache_mb(cache: Dict) -> float:
    """Estimate memory size of track_proximity_cache (numpy arrays)."""
    if not cache:
        return 0.0

    total = sys.getsizeof(cache)
    for net_id, arr in cache.items():
        total += sys.getsizeof(net_id)
        if hasattr(arr, 'nbytes'):
            # numpy array: 4 columns * 4 bytes (int32) = 16 bytes per row
            total += arr.nbytes + 128  # array overhead
        else:
            # Fallback for old dict format
            total += sys.getsizeof(arr)
    return total / (1024 * 1024)


def estimate_routed_paths_mb(paths: Dict) -> float:
    """Estimate memory size of routed_net_paths (Dict[int, List[Tuple]])."""
    if not paths:
        return 0.0

    total = sys.getsizeof(paths)
    for net_id, path in paths.items():
        total += sys.getsizeof(net_id) + sys.getsizeof(path)
        if path:
            total += len(path) * 24  # tuple of 3 ints per coordinate
    return total / (1024 * 1024)


def format_obstacle_map_stats(obstacles) -> str:
    """Format Rust GridObstacleMap statistics."""
    if obstacles is None or not hasattr(obstacles, 'get_stats'):
        return "[MEMORY] Obstacle map: N/A (no get_stats method)"

    stats = obstacles.get_stats()
    # get_stats() returns 7 counts (free_vias was added later); tolerate 6 or 7.
    blocked_cells, blocked_vias, stub_prox, layer_prox, cross_layer, source_target = stats[:6]
    free_vias = stats[6] if len(stats) > 6 else 0

    # hashbrown/FxHashMap effective bytes/entry (key+value+control, ~0.875 load):
    # (u64,u16)/(u64,i32) maps ~24B; FxHashSet<u64> ~16B. Round to ~24B for maps,
    # ~16B for the sets (allowed/source_target/free_vias are sets).
    MAP_B, SET_B = 24, 16
    map_entries = blocked_cells + blocked_vias + stub_prox + layer_prox + cross_layer
    set_entries = source_target + free_vias
    estimated_mb = (map_entries * MAP_B + set_entries * SET_B) / (1024 * 1024)

    return (f"[MEMORY] Rust obstacle map: ~{estimated_mb:.1f} MB estimated\n"
            f"         blocked_cells: {blocked_cells:,}, blocked_vias: {blocked_vias:,}\n"
            f"         stub_proximity: {stub_prox:,}, layer_proximity: {layer_prox:,}\n"
            f"         cross_layer: {cross_layer:,}, source_target: {source_target:,}, free_vias: {free_vias:,}")