#!/usr/bin/env python3
"""Guard _PyObject_DebugMallocStats C89 declarations and evaluation order."""

import re
import sys
from pathlib import Path

source = Path(sys.argv[1]).read_text()
match = re.search(
    r"\n_PyObject_DebugMallocStats\(FILE \*out\)\s*\{(.*?)\n\}",
    source, re.S,
)
assert match, "_PyObject_DebugMallocStats definition not found"
body = re.sub(r"/\*.*?\*/", "", match.group(1), flags=re.S)
declarations = r"""
    uint i;
    uint numclasses;
    size_t numpools[SMALL_REQUEST_THRESHOLD >> ALIGNMENT_SHIFT];
    size_t numblocks[SMALL_REQUEST_THRESHOLD >> ALIGNMENT_SHIFT];
    size_t numfreeblocks[SMALL_REQUEST_THRESHOLD >> ALIGNMENT_SHIFT];
    size_t allocated_bytes;
    size_t available_bytes;
    uint numfreepools;
    size_t arena_alignment;
    size_t pool_header_bytes;
    size_t quantization;
    size_t narenas;
    size_t total;
    char buf[128];
"""
assert body.split()[:len(declarations.split())] == declarations.split(), (
    "_PyObject_DebugMallocStats must hoist only uninitialized declarations "
    "before the pymalloc guard, preserving array extents"
)

# Freeze the guard before all scalar assignments, their original order, and
# the entire remaining body, including array zeroing and conditional paths.
expected = declarations + r"""
    if (!_PyMem_PymallocEnabled()) {
        return 0;
    }
    numclasses = SMALL_REQUEST_THRESHOLD >> ALIGNMENT_SHIFT;
    allocated_bytes = 0;
    available_bytes = 0;
    numfreepools = 0;
    arena_alignment = 0;
    pool_header_bytes = 0;
    quantization = 0;
    narenas = 0;
    fprintf(out, "Small block threshold = %d, in %u size classes.\n",
            SMALL_REQUEST_THRESHOLD, numclasses);

    for (i = 0; i < numclasses; ++i)
        numpools[i] = numblocks[i] = numfreeblocks[i] = 0;


    for (i = 0; i < maxarenas; ++i) {
        uintptr_t base = arenas[i].address;


        if (arenas[i].address == (uintptr_t)NULL)
            continue;
        narenas += 1;

        numfreepools += arenas[i].nfreepools;


        if (base & (uintptr_t)POOL_SIZE_MASK) {
            arena_alignment += POOL_SIZE;
            base &= ~(uintptr_t)POOL_SIZE_MASK;
            base += POOL_SIZE;
        }


        assert(base <= (uintptr_t) arenas[i].pool_address);
        for (; base < (uintptr_t) arenas[i].pool_address; base += POOL_SIZE) {
            poolp p = (poolp)base;
            const uint sz = p->szidx;
            uint freeblocks;

            if (p->ref.count == 0) {

#ifdef Py_DEBUG
                assert(pool_is_in_list(p, arenas[i].freepools));
#endif
                continue;
            }
            ++numpools[sz];
            numblocks[sz] += p->ref.count;
            freeblocks = NUMBLOCKS(sz) - p->ref.count;
            numfreeblocks[sz] += freeblocks;
#ifdef Py_DEBUG
            if (freeblocks > 0)
                assert(pool_is_in_list(p, usedpools[sz + sz]));
#endif
        }
    }
    assert(narenas == narenas_currently_allocated);

    fputc('\n', out);
    fputs("class   size   num pools   blocks in use  avail blocks\n"
          "-----   ----   ---------   -------------  ------------\n",
          out);

    for (i = 0; i < numclasses; ++i) {
        size_t p = numpools[i];
        size_t b = numblocks[i];
        size_t f = numfreeblocks[i];
        uint size = INDEX2SIZE(i);
        if (p == 0) {
            assert(b == 0 && f == 0);
            continue;
        }
        fprintf(out, "%5u %6u %11zu %15zu %13zu\n",
                i, size, p, b, f);
        allocated_bytes += b * size;
        available_bytes += f * size;
        pool_header_bytes += p * POOL_OVERHEAD;
        quantization += p * ((POOL_SIZE - POOL_OVERHEAD) % size);
    }
    fputc('\n', out);
#ifdef PYMEM_DEBUG_SERIALNO
    if (_PyMem_DebugEnabled()) {
        (void)printone(out, "# times object malloc called", serialno);
    }
#endif
    (void)printone(out, "# arenas allocated total", ntimes_arena_allocated);
    (void)printone(out, "# arenas reclaimed", ntimes_arena_allocated - narenas);
    (void)printone(out, "# arenas highwater mark", narenas_highwater);
    (void)printone(out, "# arenas allocated current", narenas);

    PyOS_snprintf(buf, sizeof(buf),
                  "%zu arenas * %d bytes/arena",
                  narenas, ARENA_SIZE);
    (void)printone(out, buf, narenas * ARENA_SIZE);

    fputc('\n', out);


    total = printone(out, "# bytes in allocated blocks", allocated_bytes);
    total += printone(out, "# bytes in available blocks", available_bytes);

    PyOS_snprintf(buf, sizeof(buf),
        "%u unused pools * %d bytes", numfreepools, POOL_SIZE);
    total += printone(out, buf, (size_t)numfreepools * POOL_SIZE);

    total += printone(out, "# bytes lost to pool headers", pool_header_bytes);
    total += printone(out, "# bytes lost to quantization", quantization);
    total += printone(out, "# bytes lost to arena alignment", arena_alignment);
    (void)printone(out, "Total", total);
    assert(narenas * ARENA_SIZE == total);

#if WITH_PYMALLOC_RADIX_TREE
    fputs("\narena map counts\n", out);
#ifdef USE_INTERIOR_NODES
    (void)printone(out, "# arena map mid nodes", arena_map_mid_count);
    (void)printone(out, "# arena map bot nodes", arena_map_bot_count);
    fputc('\n', out);
#endif
    total = printone(out, "# bytes lost to arena map root", sizeof(arena_map_root));
#ifdef USE_INTERIOR_NODES
    total += printone(out, "# bytes lost to arena map mid",
                      sizeof(arena_map_mid_t) * arena_map_mid_count);
    total += printone(out, "# bytes lost to arena map bot",
                      sizeof(arena_map_bot_t) * arena_map_bot_count);
    (void)printone(out, "Total", total);
#endif
#endif

    return 1;"""
assert body.split() == expected.split(), (
    "preserve _PyObject_DebugMallocStats early return, initialization/evaluation "
    "points, array contents, and complete operation order"
)
print("obmalloc _PyObject_DebugMallocStats C89 declarations and operation order: PASS")