TransPyC/includes/mpool.py

import t, c
from stdint import *
import viperio
import string

MPOOL_ALIGN: t.CDefine = 8
MPOOL_TYPE_SLAB: t.CDefine = 0
MPOOL_TYPE_BUMP: t.CDefine = 2


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


class MPool:
    mtype: t.CInt
    mem: t.CVoid | t.CPtr
    mem_size: t.CSizeT
    offset: t.CSizeT
    high_water: t.CSizeT
    block_size: t.CSizeT
    block_count: t.CSizeT
    used_count: t.CSizeT
    free_list: t.CVoid | t.CPtr
    alloc_map: t.CUInt8T | t.CPtr
    alloc_map_size: t.CSizeT

    def __init__(self, mem: t.CVoid | t.CPtr, mem_size: t.CSizeT, block_size: t.CSizeT):
        self.mtype = MPOOL_TYPE_BUMP
        self.mem = None
        self.mem_size = 0
        self.offset = 0
        self.high_water = 0
        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

        if mem == None: return
        if block_size > 0:
            self._init_slab(mem, mem_size, block_size)
        else:
            self._init_bump(mem, mem_size)

    def _init_bump(self, mem: t.CVoid | t.CPtr, mem_size: t.CSizeT):
        if mem == None: return
        if mem_size < 8: return

        self.mtype = MPOOL_TYPE_BUMP
        self.mem = mem
        self.mem_size = mem_size
        self.offset = 0
        self.high_water = 0

    def _init_slab(self, mem: t.CVoid | t.CPtr, mem_size: t.CSizeT, block_size: t.CSizeT):
        if mem == None: return

        bs: t.CSizeT = _align_up(block_size, MPOOL_ALIGN)
        if bs < 16: bs = 16

        alloc_map_bytes: t.CSizeT = _align_up(256, MPOOL_ALIGN)
        if mem_size < alloc_map_bytes + bs: return

        self.mtype = MPOOL_TYPE_SLAB
        self.alloc_map = mem
        self.alloc_map_size = alloc_map_bytes
        self.mem = t.CVoid(t.CUInt64T(mem) + alloc_map_bytes, t.CPtr)
        self.mem_size = mem_size - alloc_map_bytes
        self.block_size = bs
        self.used_count = 0

        idx: t.CSizeT = 0
        while idx < alloc_map_bytes:
            self.alloc_map[idx] = 0
            idx += 1

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

        self.free_list = None
        i: t.CSizeT = 0
        while i < self.block_count:
            block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.mem) + i * bs, t.CPtr)
            c.DerefAs(block, self.free_list)
            self.free_list = block
            i += 1

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

    def __exit__(self):
        if self.mtype == MPOOL_TYPE_BUMP:
            self.offset = 0
            self.high_water = 0
        elif self.mtype == MPOOL_TYPE_SLAB:
            self._slab_reset()

    def _slab_reset(self):
        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.mem) + 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

    def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        if self.mtype == MPOOL_TYPE_SLAB:
            return self._slab_alloc()
        if self.mtype == MPOOL_TYPE_BUMP:
            return self._bump_alloc(size)
        return None

    def alloc_buf(self, capacity: t.CSizeT) -> viperio.Buf | t.CPtr:
        ptr: t.CVoid | t.CPtr = self.alloc(capacity)
        if ptr == None:
            return viperio.Buf(None, 0)
        return viperio.Buf(t.CChar(t.CUInt64T(ptr), t.CPtr), capacity)

    def free(self, ptr: t.CVoid | t.CPtr):
        if ptr == None: return
        if self.mtype == MPOOL_TYPE_SLAB:
            self._slab_free(ptr)

    def realloc(self, ptr: t.CVoid | t.CPtr, old_size: t.CSizeT, new_size: t.CSizeT) -> t.CVoid | t.CPtr:
        if ptr == None:
            return self.alloc(new_size)
        if new_size == 0:
            self.free(ptr)
            return None
        if self.mtype == MPOOL_TYPE_SLAB:
            if new_size <= self.block_size:
                return ptr
        new_ptr: t.CVoid | t.CPtr = self.alloc(new_size)
        if new_ptr == None:
            return ptr
        copy_size: t.CSizeT = old_size
        if new_size < old_size:
            copy_size = new_size
        string.memcpy(new_ptr, ptr, copy_size)
        self.free(ptr)
        return new_ptr

    def reset(self):
        if self.mtype == MPOOL_TYPE_BUMP:
            self.offset = 0
            self.high_water = 0
        elif self.mtype == MPOOL_TYPE_SLAB:
            self._slab_reset()

    def _bump_alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
        if size == 0: return None
        aligned: t.CSizeT = _align_up(size, MPOOL_ALIGN)
        if self.offset + aligned > self.mem_size: return None
        ptr: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(self.mem) + self.offset, t.CPtr)
        self.offset += aligned
        if self.offset > self.high_water:
            self.high_water = self.offset
        return ptr

    def _slab_alloc(self) -> t.CVoid | t.CPtr:
        if self.free_list == None: 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.mem)) / self.block_size
        self.alloc_map[idx / 8] = self.alloc_map[idx / 8] | t.CUInt8T(1 << (idx % 8))
        return block

    def _slab_free(self, ptr: t.CVoid | t.CPtr):
        if ptr == None: return
        p: t.CUInt64T = t.CUInt64T(ptr)
        if p < t.CUInt64T(self.mem): return
        if p >= t.CUInt64T(self.mem) + self.block_count * self.block_size: return
        if (p - t.CUInt64T(self.mem)) % self.block_size != 0: return
        idx: t.CSizeT = (p - t.CUInt64T(self.mem)) / 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
        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

# 模块级包装函数（C 风格 API，第一个参数传入 pool 指针）
def bump_create(mem: t.CVoid | t.CPtr, mem_size: t.CSizeT) -> MPool | t.CPtr:
    p: MPool | t.CPtr = MPool(mem, mem_size, 0)
    return p

def alloc(pool: MPool | t.CPtr, size: t.CSizeT) -> t.CVoid | t.CPtr:
    return pool.alloc(size)

def reset(pool: MPool | t.CPtr):
    pool.reset()