SystemSimulationApp/app/simulation/property_cache.py

"""Run-local, exact-key cache for expensive thermodynamic calculations.

The cache is deliberately bound to one simulation through ``ContextVar``.
That keeps concurrent runs isolated and releases all cached states when the
run finishes.  Keys use the original Python values with no rounding or
tolerance-based reuse that could flatten numerical residuals seen by ODE and
nonlinear solvers.
"""

from __future__ import annotations

from collections.abc import Callable, Generator
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from functools import lru_cache, wraps
import os
from typing import ParamSpec, TypeVar

from app.simulation.performance import PROFILE_MODE, record_property_cache


_P = ParamSpec("_P")
_R = TypeVar("_R")
DEFAULT_PROPERTY_CACHE_MAX_ENTRIES = 8192


def _read_cache_enabled() -> bool:
    raw_value = os.getenv("SIMULATIONAPP_PROPERTY_CACHE", "on").strip().lower()
    if raw_value in {"", "1", "true", "yes", "on"}:
        return True
    if raw_value in {"0", "false", "no", "off"}:
        return False
    raise ValueError(
        "SIMULATIONAPP_PROPERTY_CACHE must be one of: on, off, true, false, 1, 0."
    )


PROPERTY_CACHE_ENABLED = _read_cache_enabled()


@dataclass(frozen=True)
class PropertyCacheInfo:
    hits: int
    misses: int
    max_entries_per_cache: int
    cache_count: int
    current_entries: int
    evictions: int


class SimulationPropertyCache:
    """Bounded C-level LRUs owned by one simulation run."""

    def __init__(self, max_entries: int = DEFAULT_PROPERTY_CACHE_MAX_ENTRIES) -> None:
        if max_entries <= 0:
            raise ValueError("Property cache max_entries must be positive.")
        self.max_entries_per_cache = int(max_entries)
        self._functions: dict[
            tuple[str, int, Callable[..., object]],
            Callable[..., object],
        ] = {}
        self._failed_misses: dict[
            tuple[str, int, Callable[..., object]],
            int,
        ] = {}
        self._owners: dict[int, tuple[object, int]] = {}
        self._next_owner_token = 0

    def owner_token(self, owner: object) -> int:
        """Return a stable identity token and retain its owner for this run."""

        identity = id(owner)
        existing = self._owners.get(identity)
        if existing is not None and existing[0] is owner:
            return existing[1]
        self._next_owner_token += 1
        self._owners[identity] = (owner, self._next_owner_token)
        return self._next_owner_token

    def get_or_compute(
        self,
        operation: str,
        owner: object,
        function: Callable[..., _R],
        args: tuple[object, ...],
        kwargs: dict[str, object],
    ) -> _R:
        cache_key = (operation, self.owner_token(owner), function)
        cached_function = self._functions.get(cache_key)
        if cached_function is None:

            @lru_cache(maxsize=self.max_entries_per_cache, typed=True)
            def invoke(*cached_args: object, **cached_kwargs: object) -> _R:
                return function(owner, *cached_args, **cached_kwargs)

            cached_function = invoke
            self._functions[cache_key] = cached_function

        if PROFILE_MODE != "audit":
            try:
                return cached_function(*args, **kwargs)
            except Exception:
                self._failed_misses[cache_key] = (
                    self._failed_misses.get(cache_key, 0) + 1
                )
                raise

        before = cached_function.cache_info()  # type: ignore[attr-defined]
        try:
            value = cached_function(*args, **kwargs)
        except Exception:
            self._failed_misses[cache_key] = (
                self._failed_misses.get(cache_key, 0) + 1
            )
            raise
        finally:
            after = cached_function.cache_info()  # type: ignore[attr-defined]
            hit = after.hits > before.hits
            record_property_cache(operation, hit=hit)
        return value

    def info(self) -> PropertyCacheInfo:
        cache_infos = {
            key: cached.cache_info()  # type: ignore[attr-defined]
            for key, cached in self._functions.items()
        }
        return PropertyCacheInfo(
            hits=sum(info.hits for info in cache_infos.values()),
            misses=sum(info.misses for info in cache_infos.values()),
            max_entries_per_cache=self.max_entries_per_cache,
            cache_count=len(self._functions),
            current_entries=sum(info.currsize for info in cache_infos.values()),
            evictions=sum(
                max(
                    0,
                    info.misses
                    - self._failed_misses.get(key, 0)
                    - info.currsize,
                )
                for key, info in cache_infos.items()
            ),
        )


_CURRENT_PROPERTY_CACHE: ContextVar[SimulationPropertyCache | None] = ContextVar(
    "simulation_property_cache",
    default=None,
)


def current_property_cache() -> SimulationPropertyCache | None:
    return _CURRENT_PROPERTY_CACHE.get()


@contextmanager
def property_cache_run(
    *,
    max_entries: int = DEFAULT_PROPERTY_CACHE_MAX_ENTRIES,
) -> Generator[SimulationPropertyCache | None, None, None]:
    """Bind a fresh cache to one top-level simulation run.

    Nested uses reuse the existing cache so lower-level simulation helpers can
    safely opt in without replacing the cache created by the API entry point.
    """

    existing = _CURRENT_PROPERTY_CACHE.get()
    if existing is not None:
        yield existing
        return
    if not PROPERTY_CACHE_ENABLED:
        yield None
        return

    cache = SimulationPropertyCache(max_entries=max_entries)
    token = _CURRENT_PROPERTY_CACHE.set(cache)
    try:
        yield cache
    finally:
        _CURRENT_PROPERTY_CACHE.reset(token)


def cache_property_calculation(
    operation: str,
) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]:
    """Cache one pure property calculation with hashable arguments per run."""

    def decorate(function: Callable[_P, _R]) -> Callable[_P, _R]:
        if not PROPERTY_CACHE_ENABLED:
            return function

        @wraps(function)
        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
            cache = _CURRENT_PROPERTY_CACHE.get()
            if cache is None:
                return function(*args, **kwargs)
            owner = args[0] if args else function
            return cache.get_or_compute(
                operation,
                owner,
                function,
                tuple(args[1:] if args else ()),
                dict(kwargs),
            )

        return wrapper

    return decorate


def with_property_cache(function: Callable[_P, _R]) -> Callable[_P, _R]:
    """Ensure a simulation entry point has a run-local cache."""

    if not PROPERTY_CACHE_ENABLED:
        return function

    @wraps(function)
    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
        with property_cache_run():
            return function(*args, **kwargs)

    return wrapper


__all__ = [
    "DEFAULT_PROPERTY_CACHE_MAX_ENTRIES",
    "PROPERTY_CACHE_ENABLED",
    "PropertyCacheInfo",
    "SimulationPropertyCache",
    "cache_property_calculation",
    "current_property_cache",
    "property_cache_run",
    "with_property_cache",
]