375 lines
14 KiB
Python
375 lines
14 KiB
Python
import t, c
|
||
from stdint import *
|
||
import string
|
||
import atom
|
||
|
||
# 伙伴系统内存分配器
|
||
# 支持可变大小分配 (malloc/free/calloc/realloc)
|
||
# 通过二分块 (buddy splitting) 和合并 (coalescing) 管理内存
|
||
# 使用 atom 库的原子操作实现自旋锁 (线程安全)
|
||
|
||
MBUDDY_MIN_BLOCK: t.CDefine = 32 # 最小块大小 (字节)，需 >= 16 以容纳链表指针
|
||
MBUDDY_MAX_ORDERS: t.CDefine = 32 # 最大阶数数量
|
||
MBUDDY_HEADER_SIZE: t.CDefine = 8 # 块头大小 (存储阶数)
|
||
|
||
|
||
def _largest_pow2_le(val: t.CSizeT) -> t.CSizeT:
|
||
"""返回 <= val 的最大 2 的幂"""
|
||
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 = MBUDDY_MIN_BLOCK
|
||
i: t.CInt
|
||
for i in range(order):
|
||
bs = bs << 1
|
||
return bs
|
||
|
||
|
||
class MBuddy:
|
||
# 编译期元数据：声明此类在 with 上下文中提供 __mbuddy__ 字段
|
||
__provides__: list[str] = ['__mbuddy__']
|
||
mem: t.CVoid | t.CPtr # 可用内存基址
|
||
mem_size: t.CSizeT # 可用内存大小 (2 的幂)
|
||
max_order: t.CInt # 最大阶数
|
||
free_lists: t.CUInt64T | t.CPtr # 空闲链头数组 (存储在 arena 开头)
|
||
lock_val: t.CVolatile | t.CInt # 自旋锁标志 (0=未锁, 1=已锁)
|
||
|
||
def __init__(self, arena: t.CVoid | t.CPtr, arena_size: t.CSizeT):
|
||
# 初始化自旋锁
|
||
self.lock_val = 0
|
||
# 空闲链数组放在 arena 开头
|
||
fl_bytes: t.CSizeT = (MBUDDY_MAX_ORDERS + 1) * 8
|
||
self.free_lists = arena
|
||
|
||
# 初始化所有空闲链头为 0 (NULL)
|
||
i: t.CInt
|
||
for i in range(MBUDDY_MAX_ORDERS + 1):
|
||
self.free_lists[i] = 0
|
||
|
||
# 计算可用区域
|
||
if arena_size <= fl_bytes:
|
||
self.mem = None
|
||
self.mem_size = 0
|
||
self.max_order = 0
|
||
return
|
||
|
||
remaining: t.CSizeT = arena_size - fl_bytes
|
||
usable: t.CSizeT = _largest_pow2_le(remaining)
|
||
|
||
if usable < MBUDDY_MIN_BLOCK:
|
||
self.mem = None
|
||
self.mem_size = 0
|
||
self.max_order = 0
|
||
return
|
||
|
||
self.mem = t.CVoid(t.CUInt64T(arena) + fl_bytes, t.CPtr)
|
||
self.mem_size = usable
|
||
|
||
# 计算 max_order: log2(usable / MIN_BLOCK)
|
||
self.max_order = 0
|
||
bs: t.CSizeT = MBUDDY_MIN_BLOCK
|
||
while bs < usable:
|
||
bs = bs << 1
|
||
self.max_order += 1
|
||
|
||
# 将整个可用区域作为一个 max_order 阶的空闲块加入空闲链
|
||
self._fl_push(self.max_order, self.mem)
|
||
|
||
def __enter__(self) -> 'MBuddy' | t.CPtr:
|
||
return self
|
||
|
||
def __exit__(self):
|
||
self.reset()
|
||
|
||
def reset(self):
|
||
"""重置分配器: 清空所有空闲链，将整个可用区域作为单个大块归还"""
|
||
if self.mem == None:
|
||
return
|
||
# 清空所有空闲链头
|
||
i: t.CInt
|
||
for i in range(MBUDDY_MAX_ORDERS + 1):
|
||
self.free_lists[i] = 0
|
||
# 将整个可用区域作为单个 max_order 阶块加入
|
||
self._fl_push(self.max_order, self.mem)
|
||
|
||
# === 空闲链操作 ===
|
||
|
||
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:
|
||
"""在指定阶数的空闲链中查找并移除 target 块。返回 1=找到, 0=未找到"""
|
||
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 cur != None:
|
||
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
|
||
|
||
# === 伙伴系统核心操作 ===
|
||
|
||
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.mem)
|
||
bs: t.CSizeT = _block_size_at_order(order)
|
||
buddy_offset: t.CSizeT = offset ^ bs
|
||
return t.CVoid(t.CUInt64T(self.mem) + buddy_offset, t.CPtr)
|
||
|
||
def _order_for_size(self, size: t.CSizeT) -> t.CInt:
|
||
"""返回使 MIN_BLOCK << order >= size 的最小阶数"""
|
||
order: t.CInt = 0
|
||
bs: t.CSizeT = MBUDDY_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:
|
||
"""找到可用块并分裂到目标阶数。返回块指针或 None"""
|
||
# 寻找第一个有空闲块的阶数 >= to_order
|
||
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)
|
||
|
||
# 逐级分裂到 to_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:
|
||
"""验证指针是否为有效的已分配块指针。返回 1=有效, 0=无效"""
|
||
if ptr == None:
|
||
return 0
|
||
if self.mem == None:
|
||
return 0
|
||
block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MBUDDY_HEADER_SIZE, t.CPtr)
|
||
# 检查范围
|
||
if t.CUInt64T(block) < t.CUInt64T(self.mem):
|
||
return 0
|
||
if t.CUInt64T(block) >= t.CUInt64T(self.mem) + self.mem_size:
|
||
return 0
|
||
# 检查对齐 (所有块都 MIN_BLOCK 对齐)
|
||
offset: t.CSizeT = t.CUInt64T(block) - t.CUInt64T(self.mem)
|
||
if offset % MBUDDY_MIN_BLOCK != 0:
|
||
return 0
|
||
# 读取头部编码值: (order << 1) | allocated_flag
|
||
stored: t.CVoid | t.CPtr = t.CVoid(c.Deref(t.CUInt64T(block, t.CPtr)), t.CPtr)
|
||
stored_val: t.CUInt64T = t.CUInt64T(stored)
|
||
# 检查 allocated 标志位 (bit 0 = 1 表示已分配)
|
||
if (stored_val & 1) == 0:
|
||
return 0 # 未分配或已释放 (双重 free 检测)
|
||
# 提取阶数
|
||
order: t.CInt = t.CInt(stored_val >> 1)
|
||
if order < 0:
|
||
return 0
|
||
if order > self.max_order:
|
||
return 0
|
||
return 1
|
||
|
||
def _fl_count(self, order: t.CInt) -> t.CSizeT:
|
||
"""统计指定阶数空闲链中的块数"""
|
||
count: t.CSizeT = 0
|
||
cur_val: t.CUInt64T = self.free_lists[order]
|
||
while cur_val != 0:
|
||
count += 1
|
||
cur: t.CVoid | t.CPtr = t.CVoid(cur_val, t.CPtr)
|
||
cur_val = t.CUInt64T(t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr))
|
||
return count
|
||
|
||
# === 公共 API ===
|
||
|
||
def _lock(self):
|
||
"""自旋锁: 原子测试并设置 lock_val，旧值为 1 时自旋等待"""
|
||
while atom.__atomic_test_and_set(c.Addr(self.lock_val), atom.ATOMIC_ACQUIRE):
|
||
pass
|
||
|
||
def _unlock(self):
|
||
"""解锁: 原子清除 lock_val"""
|
||
atom.__atomic_clear(c.Addr(self.lock_val), atom.ATOMIC_RELEASE)
|
||
|
||
def alloc(self, size: t.CSizeT) -> t.CVoid | t.CPtr:
|
||
"""分配至少 size 字节的内存。返回用户数据指针或 None (线程安全)"""
|
||
self._lock()
|
||
result: t.CVoid | t.CPtr = None
|
||
if self.mem != None:
|
||
if size != 0:
|
||
needed: t.CSizeT = size + MBUDDY_HEADER_SIZE
|
||
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 != None:
|
||
c.DerefAs(block, t.CVoid(t.CUInt64T((order << 1) | 1), t.CPtr))
|
||
result = t.CVoid(t.CUInt64T(block) + MBUDDY_HEADER_SIZE, t.CPtr)
|
||
self._unlock()
|
||
return result
|
||
|
||
def free(self, ptr: t.CVoid | t.CPtr):
|
||
"""释放之前分配的内存。非法指针将被忽略 (线程安全)"""
|
||
self._lock()
|
||
if ptr != None:
|
||
if self._is_valid_ptr(ptr) != 0:
|
||
# 从头部读取阶数 (编码: (order << 1) | 1)
|
||
block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MBUDDY_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)
|
||
# 清除 allocated 标志，防止双重 free
|
||
c.DerefAs(block, t.CVoid(0, t.CPtr))
|
||
# 合并并加入空闲链
|
||
self._coalesce(block, order)
|
||
self._unlock()
|
||
|
||
def calloc(self, count: t.CSizeT, size: t.CSizeT) -> t.CVoid | t.CPtr:
|
||
"""分配 count*size 字节并清零"""
|
||
total: t.CSizeT = count * size
|
||
ptr: t.CVoid | t.CPtr = self.alloc(total)
|
||
if ptr != 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 == None:
|
||
return self.alloc(new_size)
|
||
if new_size == 0:
|
||
self.free(ptr)
|
||
return None
|
||
if self._is_valid_ptr(ptr) == 0:
|
||
return None # 非法指针
|
||
|
||
# 读取旧阶数 (编码: (order << 1) | 1)
|
||
block: t.CVoid | t.CPtr = t.CVoid(t.CUInt64T(ptr) - MBUDDY_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 + MBUDDY_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 == None:
|
||
return ptr # 保留旧块
|
||
|
||
# 复制旧数据
|
||
old_block_size: t.CSizeT = _block_size_at_order(old_order)
|
||
old_data_size: t.CSizeT = old_block_size - MBUDDY_HEADER_SIZE
|
||
string.memcpy(new_ptr, ptr, old_data_size)
|
||
|
||
# 释放旧块
|
||
self.free(ptr)
|
||
return new_ptr
|
||
|
||
def stats(self) -> t.CSizeT:
|
||
"""返回总可用内存大小"""
|
||
return self.mem_size
|
||
|
||
def free_count(self) -> t.CSizeT:
|
||
"""统计所有空闲链中的总块数"""
|
||
total: t.CSizeT = 0
|
||
i: t.CInt
|
||
for i in range(self.max_order + 1):
|
||
total += self._fl_count(i)
|
||
return total
|
||
|
||
def self_check(self) -> t.CInt:
|
||
"""自检: 遍历所有空闲链，验证块的范围、对齐和阶数一致性。返回错误数 (0=正常)"""
|
||
if self.mem == None:
|
||
return 0
|
||
errors: t.CInt = 0
|
||
order: t.CInt
|
||
for order in range(self.max_order + 1):
|
||
bs: t.CSizeT = _block_size_at_order(order)
|
||
max_in_order: t.CSizeT = self.mem_size / bs + 1
|
||
seen: t.CSizeT = 0
|
||
cur_val: t.CUInt64T = self.free_lists[order]
|
||
while cur_val != 0:
|
||
seen += 1
|
||
if seen > max_in_order:
|
||
errors += 1
|
||
break
|
||
cur: t.CVoid | t.CPtr = t.CVoid(cur_val, t.CPtr)
|
||
# 检查范围
|
||
if t.CUInt64T(cur) < t.CUInt64T(self.mem):
|
||
errors += 1
|
||
break
|
||
if t.CUInt64T(cur) >= t.CUInt64T(self.mem) + self.mem_size:
|
||
errors += 1
|
||
break
|
||
# 检查对齐
|
||
offset: t.CSizeT = t.CUInt64T(cur) - t.CUInt64T(self.mem)
|
||
if offset % MBUDDY_MIN_BLOCK != 0:
|
||
errors += 1
|
||
break
|
||
# 检查阶数一致性: 块偏移应是其块大小的倍数
|
||
if offset % bs != 0:
|
||
errors += 1
|
||
break
|
||
cur_val = t.CUInt64T(t.CVoid(c.Deref(t.CUInt64T(cur, t.CPtr)), t.CPtr))
|
||
return errors
|