workload.py

"""A tiny workload with a slow and a fast implementation of the same function.

Use this to generate two profiles and see a real regression/improvement in the
diff.  Both implementations count duplicate pairs in a list; the slow one is
O(n^2), the fast one is O(n).  Switch which one ``run_workload`` calls by setting
the ``PROFILE_DIFF_SLOW`` environment variable.

    PROFILE_DIFF_SLOW=1 python -m profile_diff run --callable examples.workload:run_workload ...
    PROFILE_DIFF_SLOW=0 python -m profile_diff run --callable examples.workload:run_workload ...
"""

from __future__ import annotations

import os
from collections import Counter
from typing import List


def _make_data(n: int = 2000) -> List[int]:
    # Deterministic input so profiles are comparable run to run.
    return [i % 50 for i in range(n)]


def count_duplicate_pairs_slow(data: List[int]) -> int:
    """O(n^2): compare every element against every later element.

    This also allocates a throwaway list on each outer iteration, so it shows up
    in tracemalloc as well as in cProfile.
    """
    total = 0
    for i in range(len(data)):
        seen_later = [data[j] for j in range(i + 1, len(data))]  # wasteful copy
        for value in seen_later:
            if value == data[i]:
                total += 1
    return total


def count_duplicate_pairs_fast(data: List[int]) -> int:
    """O(n): count via a frequency table, then sum pair combinations."""
    counts = Counter(data)
    return sum(c * (c - 1) // 2 for c in counts.values())


def run_workload() -> int:
    """Entry point profiled by the examples; dispatches on ``PROFILE_DIFF_SLOW``."""
    data = _make_data()
    if os.environ.get("PROFILE_DIFF_SLOW", "1") == "1":
        return count_duplicate_pairs_slow(data)
    return count_duplicate_pairs_fast(data)


if __name__ == "__main__":
    print(run_workload())