gh-160310: Decompress a single stream directly in _decompress_all() - #161737

"""Per-call time of one-shot decompress() for small inputs, against the module's decompressor
alone (created and used once, as decompress() does): the difference is what decompress() adds
in Python. Also prints the decompressor's share of the call and the throughput of decompress().

    python repro_overhead.py
"""
import bz2, gzip, lzma, sys, timeit, zlib
from compression import zstd

text = b"".join(b"%d: some text to compress, " % i for i in range(1000))


def per_call(fn):
    n, _ = timeit.Timer(fn).autorange()
    return min(timeit.repeat(fn, number=n, repeat=9)) / n * 1e6


print(f"Python {sys.version.split()[0]}, best of 9, microseconds per call\n")
print("| Input | `decompress()` | decompressor alone | difference | decompressor's share | `decompress()`, MB/s |")
print("|---|---:|---:|---:|---:|---:|")
for name, compress, decompress, decompressor, sizes in (
    # gzip's decompressor alone: zlib with the gzip header and trailer, checked in C
    ("gzip", gzip.compress, gzip.decompress, lambda: zlib.decompressobj(wbits=31), (1024, 4096)),
    ("bz2", bz2.compress, bz2.decompress, bz2.BZ2Decompressor, (1024,)),
    ("lzma", lzma.compress, lzma.decompress, lzma.LZMADecompressor, (1024,)),
    ("zstd", zstd.compress, zstd.decompress, zstd.ZstdDecompressor, (1024, 4096)),
):
    for size in sizes:
        data = compress(text[:size])
        assert decompress(data) == decompressor().decompress(data) == text[:size]
        a = per_call(lambda: decompress(data))
        b = per_call(lambda: decompressor().decompress(data))
        print(f"| {name}, {size // 1024} KiB | {a:.2f} | {b:.2f} | {a - b:+.2f} | {b / a:.0%} | {size / a:,.0f} |")

"""One-shot decompress() before gh-160310, on main and with the fix, all three in one process:
the per-call time and what it adds over the decompressor alone, for the inputs of
repro_overhead.py. In each pass the states take turns, in rotating order; best of PASSES
passes, each best of 9. One process for all three takes out differences between processes: on
Windows, the bz2 and lzma decompressors alone ran up to 20% slower in some processes than in
others, which variants_overhead.py (one process per state) can't separate from the change.

    python inprocess_overhead.py PARENT_LIB FIX_STREAMS_PY [PASSES [MAIN_LIB]]

Run it without PYTHONPATH, so that the plain imports are main's; if the build's own files
aren't main's, MAIN_LIB holds main's gzip.py, bz2.py, lzma.py and compression/. PARENT_LIB holds the gzip.py,
bz2.py, lzma.py and compression/ of c3fe3c3b85 (the parent of e80962a4bd, gh-160310's fix), as
used for PYTHONPATH with variants_overhead.py; FIX_STREAMS_PY is compression/_common/_streams.py
with the fix. The parent's modules, and copies of main's bz2, lzma and zstd modules that use the
fixed _streams, are loaded under other names. gzip doesn't use _streams._decompress_all(), so
main and main + fix run the same gzip.decompress(). PASSES=0 only checks the setup.
"""
import bz2, gzip, importlib.util, lzma, os, sys, timeit, zlib
import _zstd
from compression import zstd

parent_lib, fix_path = sys.argv[1], sys.argv[2]
passes = int(sys.argv[3]) if len(sys.argv) > 3 else 7
main_lib = sys.argv[4] if len(sys.argv) > 4 else None
assert not os.environ.get("PYTHONPATH"), "run without PYTHONPATH"


def load(name, path):
    spec = importlib.util.spec_from_file_location(name, path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


streams_fix = load("_streams_fix", fix_path)
assert "first" in streams_fix._decompress_all.__code__.co_varnames, f"no fast path in {fix_path}"
if main_lib:
    streams_main = load("_streams_main",
                        os.path.join(main_lib, "compression", "_common", "_streams.py"))
else:
    from compression._common import _streams as streams_main
assert "first" not in streams_main._decompress_all.__code__.co_varnames, \
    "main's _streams already has a fast path: pass main's files as MAIN_LIB"


def with_streams(name, path, streams):
    module = load(name, path)
    module._streams = streams
    return module.decompress


def with_fix(module):
    return with_streams(module.__name__.replace(".", "_") + "_fix", module.__file__, streams_fix)


STATES = {
    "parent": {
        "gzip": load("gzip_parent", os.path.join(parent_lib, "gzip.py")).decompress,
        "bz2": load("bz2_parent", os.path.join(parent_lib, "bz2.py")).decompress,
        "lzma": load("lzma_parent", os.path.join(parent_lib, "lzma.py")).decompress,
        "zstd": load("zstd_parent",
                     os.path.join(parent_lib, "compression", "zstd", "__init__.py")).decompress,
    },
    "main": {"gzip": gzip.decompress, "bz2": bz2.decompress, "lzma": lzma.decompress,
             "zstd": zstd.decompress},
    "main + fix": {"gzip": gzip.decompress, "bz2": with_fix(bz2), "lzma": with_fix(lzma),
                   "zstd": with_fix(zstd)},
}
if main_lib:
    files = {"bz2": os.path.join(main_lib, "bz2.py"), "lzma": os.path.join(main_lib, "lzma.py"),
             "zstd": os.path.join(main_lib, "compression", "zstd", "__init__.py")}
    main_gzip = load("gzip_main", os.path.join(main_lib, "gzip.py")).decompress
    for state, suffix, streams in (("main", "_main", streams_main),
                                   ("main + fix", "_fix", streams_fix)):
        STATES[state] = {"gzip": main_gzip, **{n: with_streams(n + suffix, f, streams)
                                               for n, f in files.items()}}
# Loading compression.zstd registers its parameter enums with _zstd; give main's back.
_zstd.set_parameter_types(zstd.CompressionParameter, zstd.DecompressionParameter)
for name in ("bz2", "lzma", "zstd"):
    assert "_decompress_all" not in STATES["parent"][name].__code__.co_names, name
    assert "_decompress_all" in STATES["main"][name].__code__.co_names, name
    assert STATES["main"][name].__globals__["_streams"] is streams_main, name
    assert STATES["main + fix"][name].__globals__["_streams"] is streams_fix, name
assert STATES["parent"]["gzip"].__code__ is not STATES["main"]["gzip"].__code__

text = b"".join(b"%d: some text to compress, " % i for i in range(1000))
INPUTS = []
for name, compress, decompressor, sizes in (
    # gzip's decompressor alone: zlib with the gzip header and trailer, checked in C
    ("gzip", gzip.compress, lambda: zlib.decompressobj(wbits=31), (1024, 4096)),
    ("bz2", bz2.compress, bz2.BZ2Decompressor, (1024,)),
    ("lzma", lzma.compress, lzma.LZMADecompressor, (1024,)),
    ("zstd", zstd.compress, zstd.ZstdDecompressor, (1024, 4096)),
):
    for size in sizes:
        data = compress(text[:size])
        assert decompressor().decompress(data) == text[:size]
        for state in STATES:
            assert STATES[state][name](data) == text[:size], (name, state)
        INPUTS.append((f"{name}, {size // 1024} KiB", name, data, decompressor))


def per_call(fn):
    n, _ = timeit.Timer(fn).autorange()
    return min(timeit.repeat(fn, number=n, repeat=9)) / n * 1e6


names = list(STATES)
best = {}
for p in range(passes):
    order = names[p % 3:] + names[:p % 3]
    for row, name, data, decompressor in INPUTS:
        for state in order:
            decompress = STATES[state][name]
            a = per_call(lambda: decompress(data))
            b = per_call(lambda: decompressor().decompress(data))
            old = best.get((row, state), (float("inf"), float("inf")))
            best[row, state] = (min(old[0], a), min(old[1], b))
    print(f"pass {p + 1} of {passes} done", file=sys.stderr, flush=True)

print(f"# In one process: Python {sys.version.split()[0]}, {sys.platform}; parent = {parent_lib}, "
      f"fix = {fix_path}")
for title, value in (("decompress(), us per call", lambda a, b: f"{a:.2f}"),
                     ("added over the decompressor alone, us per call", lambda a, b: f"{a - b:+.2f}"),
                     ("the decompressor's share of the call", lambda a, b: f"{b / a:.0%}")):
    print(f"\n{title} (best of {passes} passes in one process, states taking turns, each best of 9)\n")
    print("| Input | " + " | ".join(names) + " |")
    print("|---|" + "---:|" * len(names))
    for row, *_ in INPUTS:
        if (row, names[0]) in best:
            print(f"| {row} | " + " | ".join(value(*best[row, s]) for s in names) + " |")