#!/usr/bin/env python3
"""Guard only arena_map_mark_used's entry locals and original operation order."""

import re
import sys
from pathlib import Path

source = Path(sys.argv[1]).read_text()
match = re.search(
    r"\narena_map_mark_used\(uintptr_t arena_base, int is_used\)\s*\{(.*?)\n\}",
    source, re.S,
)
assert match, "arena_map_mark_used definition not found"
body = re.sub(r"/\*.*?\*/", "", match.group(1), flags=re.S)
declarations = """
    arena_map_bot_t *n_hi;
    int i3;
    int32_t tail;
    uintptr_t arena_base_next;
    arena_map_bot_t *n_lo;
    int i3_next;
"""
assert body.split()[:len(declarations.split())] == declarations.split(), (
    "n_hi, i3, tail, arena_base_next, n_lo, i3_next must be declared without initializers at function entry"
)

# Keep evaluations at their original positions: the sanity assertion precedes
# lookup; failed lookup returns before index/tail evaluation. Lock both branches,
# including rollback and i3_next evaluation after the n_lo NULL return.
expected = declarations + """
    assert(HIGH_BITS(arena_base) == HIGH_BITS(&arena_map_root));
    n_hi = arena_map_get((block *)arena_base, is_used);
    if (n_hi == NULL) {
        assert(is_used);
        return 0;
    }
    i3 = MAP_BOT_INDEX((block *)arena_base);
    tail = (int32_t)(arena_base & ARENA_SIZE_MASK);
    if (tail == 0) {
        n_hi->arenas[i3].tail_hi = is_used ? -1 : 0;
    }
    else {
        n_hi->arenas[i3].tail_hi = is_used ? tail : 0;
        arena_base_next = arena_base + ARENA_SIZE;
        assert(arena_base < arena_base_next);
        n_lo = arena_map_get((block *)arena_base_next, is_used);
        if (n_lo == NULL) {
            assert(is_used);
            n_hi->arenas[i3].tail_hi = 0;
            return 0;
        }
        i3_next = MAP_BOT_INDEX(arena_base_next);
        n_lo->arenas[i3_next].tail_lo = is_used ? tail : 0;
    }
    return 1;
"""
assert body.split() == expected.split(), (
    "preserve original initialization/evaluation/control-flow order"
)
print("obmalloc arena_map_mark_used entry declarations and original operation order: PASS")