#!/usr/bin/env python3
"""Guard pymalloc_pool_extend's C89 declaration and operation order."""

import re
import sys
from pathlib import Path

source = Path(sys.argv[1]).read_text()
match = re.search(
    r"\npymalloc_pool_extend\(poolp pool, uint size\)\s*\{(.*?)\n\}",
    source, re.S,
)
assert match, "pymalloc_pool_extend definition not found"
body = re.sub(r"/\*.*?\*/", "", match.group(1), flags=re.S)
assert body.split()[:2] == ["poolp", "next;"], (
    "pymalloc_pool_extend must declare next without an initializer at function entry"
)

# The extension path returns before reading nextpool. Keep the full-pool
# unlink reads and writes in their original order.
expected = """
    poolp next;
    if (UNLIKELY(pool->nextoffset <= pool->maxnextoffset)) {
        pool->freeblock = (block*)pool + pool->nextoffset;
        pool->nextoffset += INDEX2SIZE(size);
        *(block **)(pool->freeblock) = NULL;
        return;
    }
    next = pool->nextpool;
    pool = pool->prevpool;
    next->prevpool = pool;
    pool->nextpool = next;
"""
assert body.split() == expected.split(), (
    "preserve pymalloc_pool_extend extension, early return, and unlink evaluation order"
)
print("obmalloc pymalloc_pool_extend C89 declaration and original operation order: PASS")


match = re.search(
    r"\nallocate_from_new_pool\(uint size\)\s*\{(.*?)\n\}",
    source, re.S,
)
assert match, "allocate_from_new_pool definition not found"
body = re.sub(r"/\*.*?\*/", "", match.group(1), flags=re.S)
assert body.split()[:2] == ["poolp", "pool;"], (
    "allocate_from_new_pool must declare pool without an initializer at function entry"
)

assert body.split()[:6] == ["poolp", "pool;", "block", "*bp;", "poolp", "next;"], (
    "allocate_from_new_pool must declare bp and next without initializers at function entry"
)
assert re.findall(r"\bblock\s+\*bp\s*;", body) == ["block *bp;"], (
    "allocate_from_new_pool must declare bp exactly once"
)
assert re.findall(r"\bpoolp\s+next\b", body) == ["poolp next"], (
    "allocate_from_new_pool must declare next exactly once"
)

# Keep the freepools read after arena creation and arena-list updates, and
# preserve all subsequent cached-pool, new-pool, and initialization branches.
# Keep next's usedpools read immediately before the frontlink writes.
expected = """
    poolp pool;
    block *bp;
    poolp next;
    if (UNLIKELY(usable_arenas == NULL)) {
#ifdef WITH_MEMORY_LIMITS
        if (narenas_currently_allocated >= MAX_ARENAS) {
            return NULL;
        }
#endif
        usable_arenas = new_arena();
        if (usable_arenas == NULL) {
            return NULL;
        }
        usable_arenas->nextarena = usable_arenas->prevarena = NULL;
        assert(nfp2lasta[usable_arenas->nfreepools] == NULL);
        nfp2lasta[usable_arenas->nfreepools] = usable_arenas;
    }
    assert(usable_arenas->address != 0);
    assert(usable_arenas->nfreepools > 0);
    if (nfp2lasta[usable_arenas->nfreepools] == usable_arenas) {
        nfp2lasta[usable_arenas->nfreepools] = NULL;
    }
    if (usable_arenas->nfreepools > 1) {
        assert(nfp2lasta[usable_arenas->nfreepools - 1] == NULL);
        nfp2lasta[usable_arenas->nfreepools - 1] = usable_arenas;
    }
    pool = usable_arenas->freepools;
    if (LIKELY(pool != NULL)) {
        usable_arenas->freepools = pool->nextpool;
        usable_arenas->nfreepools--;
        if (UNLIKELY(usable_arenas->nfreepools == 0)) {
            assert(usable_arenas->freepools == NULL);
            assert(usable_arenas->nextarena == NULL ||
                   usable_arenas->nextarena->prevarena ==
                   usable_arenas);
            usable_arenas = usable_arenas->nextarena;
            if (usable_arenas != NULL) {
                usable_arenas->prevarena = NULL;
                assert(usable_arenas->address != 0);
            }
        }
        else {
            assert(usable_arenas->freepools != NULL ||
                   usable_arenas->pool_address <=
                   (block*)usable_arenas->address +
                       ARENA_SIZE - POOL_SIZE);
        }
    }
    else {
        assert(usable_arenas->nfreepools > 0);
        assert(usable_arenas->freepools == NULL);
        pool = (poolp)usable_arenas->pool_address;
        assert((block*)pool <= (block*)usable_arenas->address +
                                 ARENA_SIZE - POOL_SIZE);
        pool->arenaindex = (uint)(usable_arenas - arenas);
        assert(&arenas[pool->arenaindex] == usable_arenas);
        pool->szidx = DUMMY_SIZE_IDX;
        usable_arenas->pool_address += POOL_SIZE;
        --usable_arenas->nfreepools;
        if (usable_arenas->nfreepools == 0) {
            assert(usable_arenas->nextarena == NULL ||
                   usable_arenas->nextarena->prevarena ==
                   usable_arenas);
            usable_arenas = usable_arenas->nextarena;
            if (usable_arenas != NULL) {
                usable_arenas->prevarena = NULL;
                assert(usable_arenas->address != 0);
            }
        }
    }
    next = usedpools[size + size];
    pool->nextpool = next;
    pool->prevpool = next;
    next->nextpool = pool;
    next->prevpool = pool;
    pool->ref.count = 1;
    if (pool->szidx == size) {
        bp = pool->freeblock;
        assert(bp != NULL);
        pool->freeblock = *(block **)bp;
        return bp;
    }
    pool->szidx = size;
    size = INDEX2SIZE(size);
    bp = (block *)pool + POOL_OVERHEAD;
    pool->nextoffset = POOL_OVERHEAD + (size << 1);
    pool->maxnextoffset = POOL_SIZE - size;
    pool->freeblock = bp + size;
    *(block **)(pool->freeblock) = NULL;
    return bp;
"""
assert body.split() == expected.split(), (
    "preserve allocate_from_new_pool arena updates, freepools assignment, and control-flow order"
)
print("obmalloc allocate_from_new_pool pool/bp/next declarations and original operation order: PASS")


match = re.search(
    r"\npymalloc_alloc\(void \*ctx, size_t nbytes\)\s*\{(.*?)\n\}",
    source, re.S,
)
assert match, "pymalloc_alloc definition not found"
body = re.sub(r"/\*.*?\*/|//[^\n]*", "", match.group(1), flags=re.S)
assert body.split()[:6] == ["uint", "size;", "poolp", "pool;", "block", "*bp;"], (
    "pymalloc_alloc must declare size, pool, and bp without initializers at function entry"
)

# Preserve the conditional Valgrind checks and request-size early returns before
# evaluating size and reading usedpools, plus the exact bp branch operations.
expected = """
    uint size;
    poolp pool;
    block *bp;
#ifdef WITH_VALGRIND
    if (UNLIKELY(running_on_valgrind == -1)) {
        running_on_valgrind = RUNNING_ON_VALGRIND;
    }
    if (UNLIKELY(running_on_valgrind)) {
        return NULL;
    }
#endif
    if (UNLIKELY(nbytes == 0)) {
        return NULL;
    }
    if (UNLIKELY(nbytes > SMALL_REQUEST_THRESHOLD)) {
        return NULL;
    }
    size = (uint)(nbytes - 1) >> ALIGNMENT_SHIFT;
    pool = usedpools[size + size];
    if (LIKELY(pool != pool->nextpool)) {
        ++pool->ref.count;
        bp = pool->freeblock;
        assert(bp != NULL);
        if (UNLIKELY((pool->freeblock = *(block **)bp) == NULL)) {
            pymalloc_pool_extend(pool, size);
        }
    }
    else {
        bp = allocate_from_new_pool(size);
    }
    return (void *)bp;
"""
assert body.split() == expected.split(), (
    "preserve pymalloc_alloc early returns, size/pool evaluation, and bp control-flow order"
)
print("obmalloc pymalloc_alloc size/pool/bp declarations and original operation order: PASS")