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()