#!/usr/bin/env python3
"""Guard _PyMem_DebugRawRealloc's C89 locals and complete operation order."""

import re
import sys
from pathlib import Path

source = Path(sys.argv[1]).read_text()
match = re.search(
    r"\n_PyMem_DebugRawRealloc\(void \*ctx, void \*p, size_t nbytes\)\s*\{(.*?)\n\}",
    source, re.S,
)
assert match, "_PyMem_DebugRawRealloc definition not found"
body = re.sub(r"/\*.*?\*/", "", match.group(1), flags=re.S)
declarations = """
    debug_alloc_api_t *api;
    uint8_t *head;
    uint8_t *data;
    uint8_t *r;
    uint8_t *tail;
    size_t total;
    size_t original_nbytes;
#define ERASED_SIZE 64
    uint8_t save[2*ERASED_SIZE];
#ifdef PYMEM_DEBUG_SERIALNO
    size_t block_serialno;
#endif
"""
assert body.split()[:len(declarations.split())] == declarations.split(), (
    "_PyMem_DebugRawRealloc must declare locals without initializers at function entry, "
    "including block_serialno under PYMEM_DEBUG_SERIALNO"
)

# Keep the ctx cast after the NULL return and the serial-number read after
# the overflow guard and tail calculation. Freeze both preprocessor paths,
# allocation failure recovery, saved-byte restoration, and growth filling.
expected = declarations + """

    if (p == NULL) {
        return _PyMem_DebugRawAlloc(0, ctx, nbytes);
    }

    api = (debug_alloc_api_t *)ctx;
    _PyMem_DebugCheckAddress(__func__, api->api_id, p);

    data = (uint8_t *)p;
    head = data - 2*SST;
    original_nbytes = read_size_t(head);
    if (nbytes > (size_t)PY_SSIZE_T_MAX - PYMEM_DEBUG_EXTRA_BYTES) {

        return NULL;
    }
    total = nbytes + PYMEM_DEBUG_EXTRA_BYTES;

    tail = data + original_nbytes;
#ifdef PYMEM_DEBUG_SERIALNO
    block_serialno = read_size_t(tail + SST);
#endif

    if (original_nbytes <= sizeof(save)) {
        memcpy(save, data, original_nbytes);
        memset(data - 2 * SST, PYMEM_DEADBYTE,
               original_nbytes + PYMEM_DEBUG_EXTRA_BYTES);
    }
    else {
        memcpy(save, data, ERASED_SIZE);
        memset(head, PYMEM_DEADBYTE, ERASED_SIZE + 2 * SST);
        memcpy(&save[ERASED_SIZE], tail - ERASED_SIZE, ERASED_SIZE);
        memset(tail - ERASED_SIZE, PYMEM_DEADBYTE,
               ERASED_SIZE + PYMEM_DEBUG_EXTRA_BYTES - 2 * SST);
    }


    r = (uint8_t *)api->alloc.realloc(api->alloc.ctx, head, total);
    if (r == NULL) {

        nbytes = original_nbytes;
    }
    else {
        head = r;
#ifdef PYMEM_DEBUG_SERIALNO
        bumpserialno();
        block_serialno = serialno;
#endif
    }
    data = head + 2*SST;

    write_size_t(head, nbytes);
    head[SST] = (uint8_t)api->api_id;
    memset(head + SST + 1, PYMEM_FORBIDDENBYTE, SST-1);

    tail = data + nbytes;
    memset(tail, PYMEM_FORBIDDENBYTE, SST);
#ifdef PYMEM_DEBUG_SERIALNO
    write_size_t(tail + SST, block_serialno);
#endif


    if (original_nbytes <= sizeof(save)) {
        memcpy(data, save, Py_MIN(nbytes, original_nbytes));
    }
    else {
        size_t i = original_nbytes - ERASED_SIZE;
        memcpy(data, save, Py_MIN(nbytes, ERASED_SIZE));
        if (nbytes > i) {
            memcpy(data + i, &save[ERASED_SIZE],
                   Py_MIN(nbytes - i, ERASED_SIZE));
        }
    }

    if (r == NULL) {
        return NULL;
    }

    if (nbytes > original_nbytes) {

        memset(data + original_nbytes, PYMEM_CLEANBYTE,
               nbytes - original_nbytes);
    }

    return data;"""
assert body.split() == expected.split(), (
    "preserve _PyMem_DebugRawRealloc initialization evaluation points, guards, "
    "and complete operation/control-flow order"
)
print("obmalloc _PyMem_DebugRawRealloc C89 declarations and original operation order: PASS")