1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
|
# SPDX-License-Identifier: MIT
from contextlib import contextmanager
__all__ = ["Heap"]
class Heap(object):
def __init__(self, start, end, block=64):
if start%block:
raise ValueError("heap start not aligned")
if end%block:
raise ValueError("heap end not aligned")
self.offset = start
self.count = (end - start) // block
self.blocks = [(self.count,False)]
self.block = block
def malloc(self, size):
size = (size + self.block - 1) // self.block
pos = 0
for i, (bsize, full) in enumerate(self.blocks):
if not full and bsize >= size:
self.blocks[i] = (size, True)
if bsize > size:
self.blocks.insert(i+1, (bsize - size, False))
return self.offset + self.block * pos
pos += bsize
raise Exception("Out of memory")
def memalign(self, align, size):
assert (align & (align - 1)) == 0
align = max(align, self.block) // self.block
size = (size + self.block - 1) // self.block
pos = self.offset // self.block
for i, (bsize, full) in enumerate(self.blocks):
if not full:
offset = 0
if pos % align:
offset = align - (pos % align)
if bsize >= (size + offset):
if offset:
self.blocks.insert(i, (offset, False))
i += 1
self.blocks[i] = (size, True)
if bsize > (size + offset):
self.blocks.insert(i+1, (bsize - size - offset, False))
return self.block * (pos + offset)
pos += bsize
raise Exception("Out of memory")
def free(self, addr):
if addr%self.block:
raise ValueError("free address not aligned")
if addr<self.offset:
raise ValueError("free address before heap")
addr -= self.offset
addr //= self.block
if addr>=self.count:
raise ValueError("free address after heap")
pos = 0
for i, (bsize, used) in enumerate(self.blocks):
if pos > addr:
raise ValueError("bad free address")
if pos == addr:
if used == False:
raise ValueError("block already free")
if i!=0 and self.blocks[i-1][1] == False:
bsize += self.blocks[i-1][0]
del self.blocks[i]
i -= 1
if i!=(len(self.blocks)-1) and self.blocks[i+1][1] == False:
bsize += self.blocks[i+1][0]
del self.blocks[i]
self.blocks[i] = (bsize, False)
return
pos += bsize
raise ValueError("bad free address")
def check(self):
free = 0
inuse = 0
for i, (bsize, used) in enumerate(self.blocks):
if used:
inuse += bsize
else:
free += bsize
if free + inuse != self.count:
raise Exception("Total block size is inconsistent")
print("Heap stats:")
print(" In use: %8dkB"%(inuse * self.block // 1024))
print(" Free: %8dkB"%(free * self.block // 1024))
@contextmanager
def guarded_malloc(self, size):
addr = self.malloc(size)
try:
yield addr
finally:
self.free(addr)
|
