import t, c
from stdint import *
import string
import atom
import viperio
import stdio
import stdlib

# ============================================================
# memhub - Unified memory management library
# MemManager (polymorphic base class, vtable)
#   ├─ MemPool  (arena bump allocator, no support for individual free)
#   ├─ MemSlab  (fixed-size block pool, alloc/free tracked via bitmap)
#   └─ MemBuddy (buddy system, supports alloc/free/realloc with coalescing)
# All subclasses override alloc/free/reset via vtable, invokable polymorphically through MemManager pointer
# ============================================================

MEMHUB_ALIGN: t.CDefine = 8

# MemSlab bitmap constants
MEMSLAB_MIN_BLOCK: t.CDefine = 16
MEMSLAB_BITMAP_BYTES: t.CDefine = 256

# MemBuddy buddy constants
MEMBUDDY_MIN_BLOCK: t.CDefine = 32
MEMBUDDY_MAX_ORDERS: t.CDefine = 32
MEMBUDDY_HEADER_SIZE: t.CDefine = 8
# 块尾 canary：alloc 时在块尾（用户数据区之外）写 8 字节哨兵，
# check_alive_blocks 校验——越界写穿数据区会破坏 canary，把
# "崩溃点漂移/事发地遥远"变成"确定性定位到越界写发生的文件"
MEMBUDDY_CANARY: t.CDefine = 8


def _canary_pattern() -> t.CUInt64T:
    """块尾 canary 值（0xA5A5A5A5A5A5A5A5，64 位）"""
    lo: t.CUInt64T = t.CUInt64T(0xA5A5A5A5)
    return (lo << 32) | lo

def _align_up(val: t.CSizeT, align: t.CSizeT) -> t.CSizeT:
    if align == 0: return val
    return (val + align - 1) & ~(align - 1)


def _largest_pow2_le(val: t.CSizeT) -> t.CSizeT:
    if val == 0: return 0
    p: t.CSizeT = 1
    while p * 2 <= val:
        p = p * 2
    return p


def _block_size_at_order(order: t.CInt) -> t.CSizeT:
    bs: t.CSizeT = MEMBUDDY_MIN_BLOCK
    i: t.CInt
    for i in range(order):
        bs = bs << 1
    return bs


# ============================================================
# MemManager - Polymorphic base class
# @t.CVTable explicit annotation to ensure recognition as vtable class during cross-module imports
# (HandlesImports only detects IsCVTable via decorator; automatic inference is unavailable across modules)
# ============================================================
@t.CVTable
class MemManager:
    __provides__: list[str] = ['__memmgr__']
    base: t.CVoid | t.CPtr        # Starting address of memory region
    size: t.CSizeT                # Total size of memory region

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT):
        self.base = base
        self.size = size

    # === Virtual functions (overridden by subclasses) ===

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        # Default: allocation not supported
        return None

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        # Default: no operation
        return 0

    def reset(self) -> t.CInt:
        # Default: no operation
        return 0

    # === Concrete methods (invoke virtual functions for automatic polymorphic dispatch) ===

    def calloc(self, count: t.CSizeT, size: t.CSizeT) -> t.CVoid | t.CPtr:
        total: t.CSizeT = count * size
        ptr: t.CVoid | t.CPtr = self.alloc(total)
        if ptr is not None:
            string.memset(ptr, 0, total)
        return ptr

    def realloc(self, ptr: t.CVoid | t.CPtr, new_size: t.CSizeT) -> t.CVoid | t.CPtr:
        if ptr is None:
            return self.alloc(new_size)
        if new_size == 0:
            self.free(ptr)
            return None
        new_ptr: t.CVoid | t.CPtr = self.alloc(new_size)
        if new_ptr is None:
            return ptr
        # Base class cannot track original allocation size; data will not be copied. Subclasses may override to preserve data
        self.free(ptr)
        return new_ptr

    def __enter__(self) -> 'MemManager' | t.CPtr:
        return self

    def __exit__(self):
        self.reset()

    def alloc_buf(self, capacity: t.CSizeT) -> viperio.Buf | t.CPtr:
        buf: viperio.Buf | t.CPtr = self.alloc(viperio.Buf.__sizeof__())
        if buf is None:
            return None
        data_ptr: t.CVoid | t.CPtr = self.alloc(capacity)
        if data_ptr is None:
            return None
        buf.__before_init__()
        buf.__init__(t.CChar(t.CUInt64T(data_ptr), t.CPtr), capacity)
        return buf


# ============================================================
# MemPool - Arena bump allocator
# Individual free operations are unsupported; reset reclaims all memory at once
# ============================================================
class MemPool(MemManager):
    offset: t.CSizeT           # Bump pointer offset
    high_water: t.CSizeT       # High-water mark

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT):
        self.base = base
        self.size = size
        self.offset = 0
        self.high_water = 0

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        if size == 0: return None
        aligned: t.CSizeT = _align_up(size, MEMHUB_ALIGN)
        if self.offset + aligned > self.size:
            stdio.printf("\n[OOM] 内存不足 (Out of Memory)\n")
            stdio.printf("[OOM] 尝试分配: %zu 字节，失败\n", size)
            stdio.printf("[OOM] MemPool: total=%zu used=%zu (bump allocator, 无释放)\n",
                         self.size, self.offset)
            stdio.printf("[OOM] 说明: 内存池已耗尽。可能原因: (1) 池大小不足，需增大 POOL_SIZE; (2) 存在内存泄漏。\n")
            stdlib.exit(1)
        ptr: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.base) + self.offset, t.CPtr)
        self.offset += aligned
        if self.offset > self.high_water:
            self.high_water = self.offset
        return ptr

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        # Bump allocator does not support individual free operations
        return 1

    def reset(self) -> t.CInt:
        self.offset = 0
        self.high_water = 0
        return 1


# ============================================================
# MemSlab - Fixed-size block pool
# Bitmap stored at the start of arena; subsequent area split into block_size units managed by free list
# ============================================================
class MemSlab(MemManager):
    block_size: t.CSizeT               # Block size (aligned)
    block_count: t.CSizeT              # Total block count
    used_count: t.CSizeT               # Used block count
    free_list: t.CVoid | t.CPtr        # Free list head
    alloc_map: t.CUInt8T | t.CPtr      # Allocation bitmap (at start of arena)
    alloc_map_size: t.CSizeT           # Bitmap byte count
    usable: t.CVoid | t.CPtr           # Usable block area start (after bitmap)
    usable_size: t.CSizeT              # Usable block area size

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT, block_size: t.CSizeT):
        self.base = base
        self.size = size
        self.block_size = 0
        self.block_count = 0
        self.used_count = 0
        self.free_list = None
        self.alloc_map = None
        self.alloc_map_size = 0
        self.usable = None
        self.usable_size = 0

        if base is None: return

        bs: t.CSizeT = _align_up(block_size, MEMHUB_ALIGN)
        if bs < MEMSLAB_MIN_BLOCK: bs = MEMSLAB_MIN_BLOCK

        map_bytes: t.CSizeT = _align_up(MEMSLAB_BITMAP_BYTES, MEMHUB_ALIGN)
        if size < map_bytes + bs: return

        self.alloc_map = base
        self.alloc_map_size = map_bytes
        self.usable = t.CVoid(t.CUInt64T(base) + map_bytes, t.CPtr)
        self.usable_size = size - map_bytes
        self.block_size = bs

        # Clear bitmap
        idx: t.CSizeT = 0
        while idx < map_bytes:
            self.alloc_map[idx] = 0
            idx += 1

        self.block_count = self.usable_size / bs
        if self.block_count == 0: return

        # Construct free list: the first 8 bytes of each block store pointer to next block
        self.free_list = None
        i: t.CSizeT = 0
        while i < self.block_count:
            block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.usable) + i * bs, t.CPtr)
            c.DerefAs(block, self.free_list)
            self.free_list = block
            i += 1

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        # Slab mode: allocate one block only if requested size <= block_size
        if self.free_list is None:
            stdio.printf("\n[OOM] 内存不足 (Out of Memory)\n")
            stdio.printf("[OOM] 尝试分配: %zu 字节，失败（free_list 已空）\n", size)
            stdio.printf("[OOM] MemSlab: block_size=%zu block_count=%zu used=%zu\n",
                         self.block_size, self.block_count, self.used_count)
            stdio.printf("[OOM] 说明: 内存池已耗尽。可能原因: (1) 池大小不足，需增大 POOL_SIZE; (2) 存在内存泄漏。\n")
            stdlib.exit(1)
        if size > self.block_size:
            # 超块大小请求 = 调用方错误：优雅失败返回 NULL（类 malloc 语义，调用方自行判断）。
            # 与 free_list 耗尽（真正的 OOM，打印后 exit(1)）区分开。
            stdio.printf("[WARN] MemSlab: 请求 %zu 字节超过 block_size=%zu，分配失败返回 NULL；应改用 MemBuddy/MemPool。\n",
                         size, self.block_size)
            return None
        block: t.CVoid | t.CPtr = self.free_list
        self.free_list = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        self.used_count += 1
        idx: t.CSizeT = (t.CUInt64T(block) - t.CUInt64T(self.usable)) / self.block_size
        self.alloc_map[idx / 8] = self.alloc_map[idx / 8] | t.CUInt8T(1 << (idx % 8))
        return block

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        if ptr is None: return 0
        p: t.CUInt64T = t.CUInt64T(ptr)
        if p < t.CUInt64T(self.usable): return 0
        if p >= t.CUInt64T(self.usable) + self.block_count * self.block_size: return 0
        if (p - t.CUInt64T(self.usable)) % self.block_size != 0: return 0
        idx: t.CSizeT = (p - t.CUInt64T(self.usable)) / self.block_size
        byte_idx: t.CSizeT = idx / 8
        bit_idx: t.CInt = t.CInt(idx % 8)
        if (self.alloc_map[byte_idx] >> bit_idx) & 1 == 0: return 0
        self.alloc_map[byte_idx] = self.alloc_map[byte_idx] & t.CUInt8T(~(1 << bit_idx))
        c.DerefAs(ptr, self.free_list)
        self.free_list = ptr
        self.used_count -= 1
        return 1

    def reset(self) -> t.CInt:
        self.used_count = 0
        self.free_list = None
        i: t.CSizeT = 0
        while i < self.block_count:
            block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.usable) + i * self.block_size, t.CPtr)
            c.DerefAs(block, self.free_list)
            self.free_list = block
            if i / 8 < self.alloc_map_size:
                self.alloc_map[i / 8] = 0
            i += 1
        return 1


# ============================================================
# MemBuddy - Buddy system allocator
# Free list head array stored at arena start; subsequent memory managed in power-of-two sized blocks
# Thread-safe (spinlock enabled), supports block coalescing
#
# Checkpoint 支持（供编译器等长任务回收阶段性分配）:
#   mark()      → 返回当前分配水位（不分配内存）
#   reset_to(m) → 释放水位 m 之后分配的所有块（回收到 free list）
# 记录栈从 usable 顶部预留 _CHK_RESERVE 字节，不参与 buddy 分配。
# ============================================================
class MemBuddy(MemManager):
    # Override parent __provides__: MemBuddy provides both __mbuddy__ and __memmgr__.
    # Within `with MemBuddy(...)` context, classes (e.g. _str) with __requires__=['__mbuddy__']
    # can be automatically injected via _find_provider
    __provides__: list[str] = ['__mbuddy__', '__memmgr__']
    max_order: t.CInt                  # Maximum order
    free_lists: t.CUInt64T | t.CPtr    # Free list head array (located at arena start)
    lock_val: t.CVolatile | t.CInt     # Spinlock
    usable: t.CVoid | t.CPtr           # Start of usable memory region (skips free_lists)
    usable_size: t.CSizeT              # Size of usable memory (power of two)
    chk_stack: t.CUInt64T | t.CPtr     # Checkpoint 分配记录栈（存 block 地址）
    chk_cap: t.CSizeT                  # 记录栈容量（条目数）
    chk_depth: t.CSizeT                # 当前记录深度
    chk_base: t.CVoid | t.CPtr         # 记录栈起始地址（由 stdlib.malloc 分配，独立于 buddy 区）
    big_alloc_count: t.CSizeT          # ≥8KB 大分配计数（OOM 诊断：定位大缓冲反复分配）

    def __init__(self, base: t.CVoid | t.CPtr, size: t.CSizeT):
        self.base = base
        self.size = size
        self.max_order = 0
        self.free_lists = None
        self.lock_val = 0
        self.usable = None
        self.usable_size = 0
        self.chk_stack = None
        self.chk_cap = 0
        self.chk_depth = 0
        self.chk_base = None
        self.big_alloc_count = 0

        self.lock_val = 0
        fl_bytes: t.CSizeT = (MEMBUDDY_MAX_ORDERS + 1) * 8
        self.free_lists = base

        # Initialize all free list heads to be 0
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            self.free_lists[i] = 0

        if size <= fl_bytes:
            self.usable = None
            self.usable_size = 0
            self.max_order = 0
            return

        remaining: t.CSizeT = size - fl_bytes

        # 注意：不能在 __init__ 里预留 checkpoint 记录栈！
        # 记录栈改为 mark() 时从 buddy 池懒分配（见 mark()）。
        #
        # 非 2 的幂 arena 支持：主池 1.2GB（非 2 幂）不再被
        # _largest_pow2_le 砍成 1GB 浪费 0.2GB。整个区域按二进制
        # 分解拆成若干个 2 的幂空闲块压入对应 free list；buddy 合并
        # 按 offset^bs 匹配，跨分解边界的块因 order 不匹配自然停止
        # 合并，不会破坏 buddy 一致性。remaining 恰为 2 的幂时分解
        # 结果与原实现（单个 max_order 块）完全等价。
        max_blk: t.CSizeT = _largest_pow2_le(remaining)
        if max_blk < MEMBUDDY_MIN_BLOCK:
            self.usable = None
            self.usable_size = 0
            self.max_order = 0
            return

        self.usable = t.CVoid(t.CUInt64T(base) + fl_bytes, t.CPtr)

        # 计算 max_order: log2(max_blk / MIN_BLOCK)
        self.max_order = 0
        bs: t.CSizeT = MEMBUDDY_MIN_BLOCK
        while bs < max_blk:
            bs = bs << 1
            self.max_order += 1

        # 二进制分解：把 [usable, usable+remaining) 拆成 2 的幂空闲块。
        # 尾部碎片（< MIN_BLOCK）不入 free list；usable_size 记实际
        # 可分配字节（used/free 统计不含碎片）。
        pos: t.CSizeT = 0
        rem: t.CSizeT = remaining
        used_area: t.CSizeT = 0
        while rem >= MEMBUDDY_MIN_BLOCK:
            blk: t.CSizeT = _largest_pow2_le(rem)
            o: t.CInt = 0
            ob: t.CSizeT = MEMBUDDY_MIN_BLOCK
            while ob < blk:
                ob = ob << 1
                o += 1
            self._fl_push(o, t.CVoid(t.CUInt64T(self.usable) + pos, t.CPtr))
            pos += blk
            used_area += blk
            rem -= blk
        self.usable_size = used_area

    # === Checkpoint support ===

    def mark(self) -> t.CSizeT:
        """返回当前分配水位（checkpoint）

        记录栈用 stdlib.malloc 分配（独立于 buddy 区，避免占用 buddy 池、
        避免 reset_to 误释放栈自身，也避免栈内存与 buddy 块重叠）。
        初始容量 2M 条目（16MB）：单文件翻译（parse AST + 生成 IR）可产生
        几十万次 alloc，64K 初始容量会在首个大文件就触发 realloc 扩容，
        扩容失败时块不再记录 → reset_to 永久 no-op → 所有分配泄漏 →
        1GB 主池 OOM（NewAstTest Phase2 根因）。
        """
        if self.chk_stack is None:
            cap: t.CSizeT = 2 * 1024 * 1024
            base: t.CVoid | t.CPtr = t.CVoid(stdlib.malloc(cap * 8), t.CPtr)
            if base is None:
                return 0
            self.chk_base = base
            self.chk_stack = (t.CUInt64T | t.CPtr)(t.CVoid(base, t.CPtr))
            self.chk_cap = cap
            self.chk_depth = 0
        return self.chk_depth

    def reset_to(self, mark: t.CSizeT):
        """释放水位 mark 之后分配的所有块，回收到 free list

        栈中每个 block：若 header 仍标记为已分配（奇数），_coalesce 归还；
        若已被单独 free（header==0），跳过避免重复归还。
        """
        if self.chk_stack is None:
            return
        if mark > self.chk_depth:
            # mark 大于当前深度：调用方 mark 值异常（可能读错字段/跨实例），
            # reset 将 no-op → 之后分配全泄漏 → 1GB 主池 OOM。fail-fast 暴露。
            stdio.printf(
                "[MEM] reset_to mark=%zu > chk_depth=%zu（mark 异常，reset 失效，终止）\n",
                mark, self.chk_depth)
            return
        if mark >= self.chk_depth:
            return
        self._lock()
        while self.chk_depth > mark:
            self.chk_depth -= 1
            block: t.CVoid | t.CPtr = t.CVoid(self.chk_stack[self.chk_depth], t.CPtr)
            if block is None:
                continue
            stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
            stored_val: t.CUInt64T = t.CUInt64T(stored)
            if (stored_val & 1) == 0:
                continue  # 已被单独 free，跳过
            order: t.CInt = t.CInt(stored_val >> 1)
            if order < 0 or order > self.max_order:
                continue
            self._coalesce(block, order)
        self._unlock()

    def reset_all(self):
        """回收全部已分配块（等价 reset_to(0)）"""
        if self.chk_stack is None:
            return
        self.reset_to(0)

    # === Free list operations ===

    def _fl_push(self, order: t.CInt, block: t.CVoid | t.CPtr):
        old_head: t.CUInt64T = self.free_lists[order]
        c.DerefAs(block, t.CVoid(old_head, t.CPtr))
        self.free_lists[order] = t.CUInt64T(block)

    def _fl_pop(self, order: t.CInt) -> t.CVoid | t.CPtr:
        head_val: t.CUInt64T = self.free_lists[order]
        if head_val == 0:
            return None
        block: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
        next_ptr: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        self.free_lists[order] = t.CUInt64T(next_ptr)
        return block

    def _fl_find_and_remove(self, order: t.CInt, target: t.CVoid | t.CPtr) -> t.CInt:
        head_val: t.CUInt64T = self.free_lists[order]
        if head_val == 0:
            return 0
        head: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
        if t.CUInt64T(head) == t.CUInt64T(target):
            next_ptr: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(head, t.CPtr)), t.CPtr)
            self.free_lists[order] = t.CUInt64T(next_ptr)
            return 1
        prev: t.CVoid | t.CPtr = head
        cur: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(head, t.CPtr)), t.CPtr)
        while t.CUInt64T(cur) != 0:
            if t.CUInt64T(cur) == t.CUInt64T(target):
                next_ptr: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
                c.DerefAs(prev, next_ptr)
                return 1
            prev = cur
            cur = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
        return 0

    # === Buddy system core ===

    def _buddy_of(self, block: t.CVoid | t.CPtr, order: t.CInt) -> t.CVoid | t.CPtr:
        offset: t.CSizeT = t.CUInt64T(block) - t.CUInt64T(self.usable)
        bs: t.CSizeT = _block_size_at_order(order)
        buddy_offset: t.CSizeT = offset ^ bs
        return t.CVoid(t.CUInt64T(self.usable) + buddy_offset, t.CPtr)

    def _order_for_size(self, size: t.CSizeT) -> t.CInt:
        order: t.CInt = 0
        bs: t.CSizeT = MEMBUDDY_MIN_BLOCK
        while bs < size:
            bs = bs << 1
            order += 1
        return order

    def _split_to_order(self, to_order: t.CInt) -> t.CVoid | t.CPtr:
        found_order: t.CInt = to_order
        while found_order <= self.max_order:
            if self.free_lists[found_order] != 0:
                break
            found_order += 1
        if found_order > self.max_order:
            return None
        block: t.CVoid | t.CPtr = self._fl_pop(found_order)
        while found_order > to_order:
            found_order -= 1
            bs: t.CSizeT = _block_size_at_order(found_order)
            buddy: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(block) + bs, t.CPtr)
            self._fl_push(found_order, buddy)
        return block

    def _coalesce(self, block: t.CVoid | t.CPtr, order: t.CInt):
        while order < self.max_order:
            buddy: t.CVoid | t.CPtr = self._buddy_of(block, order)
            found: t.CInt = self._fl_find_and_remove(order, buddy)
            if found == 0:
                break
            if t.CUInt64T(buddy) < t.CUInt64T(block):
                block = buddy
            order += 1
        self._fl_push(order, block)

    def _is_valid_ptr(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        if ptr is None: return 0
        if self.usable is None: return 0
        block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MEMBUDDY_HEADER_SIZE, t.CPtr)
        if t.CUInt64T(block) < t.CUInt64T(self.usable): return 0
        if t.CUInt64T(block) >= t.CUInt64T(self.usable) + self.usable_size: return 0
        offset: t.CSizeT = t.CUInt64T(block) - t.CUInt64T(self.usable)
        if offset % MEMBUDDY_MIN_BLOCK != 0: return 0
        stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        stored_val: t.CUInt64T = t.CUInt64T(stored)
        if (stored_val & 1) == 0: return 0
        order: t.CInt = t.CInt(stored_val >> 1)
        if order < 0: return 0
        if order > self.max_order: return 0
        # 记录栈一致性校验：分配时（alloc）每个 block 地址都会压入 chk_stack。
        # reset_to(mark) 会把 [mark, chk_depth) 的块弹栈并归还 free list，
        # 之后这些地址会被重新分配。若一个 header 恰好是合法奇数（假阳性），
        # 但地址不在当前 chk_stack[0..chk_depth) 存活区间内，说明已被回收，
        # 不能当作有效指针（否则 _format_func_type_str/TypePrint 会沿垃圾链
        # 无限分配导致 OOM）。
        # 注意：chk_stack 容量满（64K）后 alloc 不再记录新块（chk_depth==chk_cap），
        # 此时无法区分，跳过校验（退化到仅 header 校验，避免误判正常 free）。
        if self.chk_stack is not None and self.chk_depth < self.chk_cap:
            alive: t.CInt = 0
            ci: t.CSizeT = 0
            while ci < self.chk_depth:
                if self.chk_stack[ci] == t.CUInt64T(block):
                    alive = 1
                    break
                ci += 1
            if alive == 0:
                return 0
        return 1

    # === Spinlock ===

    def _lock(self): # can use the spinlock library for decoupling
        while atom.__atomic_test_and_set(c.Addr(self.lock_val), atom.ATOMIC_ACQUIRE):
            pass

    def _unlock(self):
        atom.__atomic_clear(c.Addr(self.lock_val), atom.ATOMIC_RELEASE)

    # === Virtual function overrides ===

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        self._lock()
        result: t.CVoid | t.CPtr = None
        if size != 0 and self.usable is not None:
            # +MEMBUDDY_CANARY：预留块尾 8 字节哨兵（用户数据区之外），越界写定位用
            needed: t.CSizeT = size + MEMBUDDY_HEADER_SIZE + MEMBUDDY_CANARY
            order: t.CInt = self._order_for_size(needed)
            if order <= self.max_order:
                block: t.CVoid | t.CPtr = self._split_to_order(order)
                if block is not None:
                    c.DerefAs(block, t.CVoid(t.CUInt64T((order << 1) | 1), t.CPtr))
                    result = t.CVoid(t.CUInt64T(block) + MEMBUDDY_HEADER_SIZE, t.CPtr)
                    # 写块尾 canary（block + block_size - CANARY）
                    bs: t.CSizeT = _block_size_at_order(order)
                    c.DerefAs(t.CVoid(t.CUInt64T(block) + bs - MEMBUDDY_CANARY, t.CPtr),
                              t.CVoid(_canary_pattern(), t.CPtr))
                    # 记录到 checkpoint 栈（供 reset_to 回收）
                    # 栈满时 stdlib.realloc 翻倍扩容（防爆满后记录缺失
                    # → reset 不回收 + _is_valid_ptr 栈校验退化的连锁问题）
                    if self.chk_stack is not None:
                        if self.chk_depth >= self.chk_cap:
                            new_cap: t.CSizeT = self.chk_cap * 2
                            new_stack: t.CVoid | t.CPtr = t.CVoid(
                                stdlib.realloc(self.chk_base, new_cap * 8), t.CPtr)
                            if new_stack is not None:
                                self.chk_base = new_stack
                                self.chk_stack = (t.CUInt64T | t.CPtr)(t.CVoid(new_stack, t.CPtr))
                                self.chk_cap = new_cap
                            else:
                                # realloc 失败：块不再记录 → reset_to 无法回收该块 →
                                # 渐进泄漏最终 1GB 主池 OOM。fail-fast 暴露而非静默。
                                stdio.printf(
                                    "[MEM] 记录栈 realloc 失败 cap=%zu depth=%zu（reset 将失效，终止）\n",
                                    self.chk_cap, self.chk_depth)
                        if self.chk_depth < self.chk_cap:
                            self.chk_stack[self.chk_depth] = t.CUInt64T(block)
                            self.chk_depth += 1
        if result is None:
            # 分配失败诊断（定位"异常 size/池耗尽导致 None → 调用方解引用段错误"）
            if size != 0 and self.usable is not None:
                needed2: t.CSizeT = size + MEMBUDDY_HEADER_SIZE + MEMBUDDY_CANARY
                order2: t.CInt = self._order_for_size(needed2)
                stdio.printf("[MEM] alloc 失败 size=%zu needed=%zu order=%d max_order=%d used=%zu\n",
                             size, needed2, order2, self.max_order, self.used_bytes())
                # 越界写诊断：free list 结构损坏（self_check!=0）说明有越界写穿
                # 破坏块头/free list next 指针 → alloc 找不到块（used 虚高非真实耗尽）
                sc: t.CInt = self.self_check()
                stdio.printf("[MEM] alloc 失败 self_check=%d free_bytes=%zu free_count=%zu chk_depth=%zu big=%zu\n",
                             sc, self.free_bytes(), self.free_count(), self.chk_depth, self.big_alloc_count)
                # chk_depth 巨大（数百万）→ 分配次数爆炸（每语句/每调用反复分配）；
                # chk_depth 不大但 used 巨大 → 单次大分配（缓冲/表扩容）反复发生
                if order2 > self.max_order:
                    stdio.printf("[MEM] 原因：请求 size 超 pool 最大块（order %d > max %d）\n",
                                 order2, self.max_order)
                else:
                    stdio.printf("[MEM] 原因：pool 耗尽（空闲 %zu 字节不足）\n", self.free_bytes())
            # 分配失败（alloc(0) 或池耗尽）：解锁并返回 None，由调用者决定后续处理。
            # 不能 exit——测试依赖 alloc 失败返回 None 后继续 free+重试
            # （如 OOM section: 耗尽后 free 一块再 alloc 应成功）。
            self._unlock()
            return None
        self._unlock()
        return result

    def free(self, ptr: t.CVoid | t.CPtr) -> t.CInt:
        self._lock()
        if ptr is not None:
            if self._is_valid_ptr(ptr) != 0:
                block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MEMBUDDY_HEADER_SIZE, t.CPtr)
                stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
                stored_val: t.CUInt64T = t.CUInt64T(stored)
                order: t.CInt = t.CInt(stored_val >> 1)
                c.DerefAs(block, t.CVoid(0, t.CPtr))
                self._coalesce(block, order)
        self._unlock()
        return 1

    def _rebuild_free_lists(self):
        """按二进制分解重建空闲链表（reset 恢复初始分配格局）

        与 __init__ 的初始化分解一致：非 2 的幂 arena 会拆成多个
        2 的幂空闲块（buddy 合并跨分解边界时因 order 不匹配自然
        停止，安全）；2 的幂 arena 等价于单个 max_order 块。
        """
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            self.free_lists[i] = 0
        pos: t.CSizeT = 0
        rem: t.CSizeT = self.usable_size
        while rem >= MEMBUDDY_MIN_BLOCK:
            blk: t.CSizeT = _largest_pow2_le(rem)
            o: t.CInt = 0
            ob: t.CSizeT = MEMBUDDY_MIN_BLOCK
            while ob < blk:
                ob = ob << 1
                o += 1
            self._fl_push(o, t.CVoid(t.CUInt64T(self.usable) + pos, t.CPtr))
            pos += blk
            rem -= blk

    def reset(self) -> t.CInt:
        if self.usable is None:
            return 0
        self._rebuild_free_lists()
        return 1

    # realloc override: buddy system reads original order from block header; old_size is not required
    def realloc(self, ptr: t.CVoid | t.CPtr, new_size: t.CSizeT) -> t.CVoid | t.CPtr:
        if ptr is None:
            return self.alloc(new_size)
        if new_size == 0:
            self.free(ptr)
            return None
        if self._is_valid_ptr(ptr) == 0:
            return None
        block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MEMBUDDY_HEADER_SIZE, t.CPtr)
        stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
        stored_val: t.CUInt64T = t.CUInt64T(stored)
        old_order: t.CInt = t.CInt(stored_val >> 1)
        needed: t.CSizeT = new_size + MEMBUDDY_HEADER_SIZE
        new_order: t.CInt = self._order_for_size(needed)
        if new_order <= old_order:
            return ptr
        new_ptr: t.CVoid | t.CPtr = self.alloc(new_size)
        if new_ptr is None:
            return ptr
        old_block_size: t.CSizeT = _block_size_at_order(old_order)
        old_data_size: t.CSizeT = old_block_size - MEMBUDDY_HEADER_SIZE
        string.memcpy(new_ptr, ptr, old_data_size)
        self.free(ptr)
        return new_ptr

    # === Status Queries ===

    @property
    def mem_size(self) -> t.CSizeT:
        # Total usable region size (power of two). Maximum allocatable size = mem_size - MEMBUDDY_HEADER_SIZE
        return self.usable_size

    def stats(self) -> t.CSizeT:
        # Return total usable region size, for status statistics
        return self.usable_size

    def free_count(self) -> t.CSizeT:
        # Count total blocks in all order lists for status statistics
        count: t.CSizeT = 0
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            head_val: t.CUInt64T = self.free_lists[i]
            cur: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
            while t.CUInt64T(cur) != 0:
                count += 1
                cur = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
        return count

    def free_bytes(self) -> t.CSizeT:
        # Count total bytes in all free blocks for status statistics
        total: t.CSizeT = 0
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            head_val: t.CUInt64T = self.free_lists[i]
            cur: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
            while t.CUInt64T(cur) != 0:
                total += _block_size_at_order(i)
                cur = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
        return total

    def used_bytes(self) -> t.CSizeT:
        # Used bytes = total usable region size - free bytes in all free blocks
        return self.usable_size - self.free_bytes()

    def dump_stats(self, label: t.CChar | t.CPtr):
        # Print memory usage statistics
        fb: t.CSizeT = self.free_bytes()
        ub: t.CSizeT = self.usable_size - fb
        stdio.printf("[MEM] %s: total=%zu used=%zu (%zu%%) free=%zu free_blocks=%zu\n",
                     label, self.usable_size, ub,
                     (ub * 100) / self.usable_size if self.usable_size > 0 else 0,
                     fb, self.free_count())

    def self_check(self) -> t.CInt:
        # Validate buddy allocator internal consistency, return 0=OK, non-0=damage
        if self.usable is None:
            return 0
        # 环检测：free list 被越界写破坏 next 指针后可能成环，
        # 无步数上限会导致无限循环（崩溃点漂移排查哨兵必须能终止）。
        max_steps: t.CSizeT = self.usable_size / MEMBUDDY_MIN_BLOCK + 2
        i: t.CInt
        for i in range(MEMBUDDY_MAX_ORDERS + 1):
            head_val: t.CUInt64T = self.free_lists[i]
            cur: t.CVoid | t.CPtr = t.CVoid(head_val, t.CPtr)
            steps: t.CSizeT = 0
            while t.CUInt64T(cur) != 0:
                steps += 1
                if steps > max_steps:
                    return 4  # free list 成环或节点数超上限（next 指针被破坏）
                # Check pointer is in usable region range
                if t.CUInt64T(cur) < t.CUInt64T(self.usable):
                    return 1
                if t.CUInt64T(cur) >= t.CUInt64T(self.usable) + self.usable_size:
                    return 2
                # Check alignment: order i blocks must be aligned to size _block_size_at_order(i)
                offset: t.CSizeT = t.CUInt64T(cur) - t.CUInt64T(self.usable)
                bs: t.CSizeT = _block_size_at_order(i)
                if offset % bs != 0:
                    return 3
                cur = t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr)
        return 0

    def check_alive_blocks(self, from_mark: t.CSizeT) -> t.CInt:
        # 检查 [from_mark, chk_depth) 记录的所有已分配块 header 完整性。
        # 越界写（参数错位/下标 GEP 越界）会破坏相邻块 header（奇偶位/order 错乱），
        # 返回 0=OK，非 0=损坏（rc 语义见调用方注释）。
        # 用于 reset_to 前哨兵：把"崩溃点漂移"变成"确定性定位到损坏发生的文件"。
        if self.chk_stack is None:
            return
        if from_mark > self.chk_depth:
            # from_mark 大于当前深度：调用方 mark 值异常（可能读错字段/跨实例），
            # 校验将 no-op → 越界写漏检。fail-fast 暴露。
            stdio.printf(
                "[MEM] check_alive_blocks from_mark=%zu > chk_depth=%zu（mark 异常）\n",
                from_mark, self.chk_depth)
            return
        if from_mark >= self.chk_depth:
            return 0
        if self.usable is None:
            return 0
        ci: t.CSizeT = from_mark
        while ci < self.chk_depth:
            block: t.CVoid | t.CPtr = t.CVoid(self.chk_stack[ci], t.CPtr)
            if block is not None:
                if t.CUInt64T(block) < t.CUInt64T(self.usable):
                    return 1
                if t.CUInt64T(block) >= t.CUInt64T(self.usable) + self.usable_size:
                    return 2
                stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
                stored_val: t.CUInt64T = t.CUInt64T(stored)
                if (stored_val & 1) != 0:
                    # 活块：header 必须是 (order<<1)|1，order 合法且块对齐到该 order 大小
                    order: t.CInt = t.CInt(stored_val >> 1)
                    if order < 0:
                        return 3
                    if order > self.max_order:
                        return 4
                    offset: t.CSizeT = t.CUInt64T(block) - t.CUInt64T(self.usable)
                    bs: t.CSizeT = _block_size_at_order(order)
                    if offset % bs != 0:
                        return 5
                    # 块尾 canary 校验：越界写穿数据区会破坏块尾哨兵，
                    # 立即定位到越界写发生的文件（崩溃点漂移的关键哨兵）
                    can_stored: t.CVoid | t.CPtr = t.CVoid(
                        c.Deref(t.CUInt64T(t.CUInt64T(block) + bs - MEMBUDDY_CANARY, t.CPtr), t.CPtr),
                        t.CPtr)
                    if t.CUInt64T(can_stored) != _canary_pattern():
                        return 6
            ci += 1
        return 0