Line data Source code
1 : //https://github.com/cloudwu/buddy
2 :
3 : #include <dash/dart/mpi/dart_mem.h>
4 :
5 : /* For PRIu64, uint64_t in printf */
6 : #define __STDC_FORMAT_MACROS
7 : #include <inttypes.h>
8 :
9 : #define NODE_UNUSED 0
10 : #define NODE_USED 1
11 : #define NODE_SPLIT 2
12 : #define NODE_FULL 3
13 :
14 : struct dart_buddy *
15 4 : dart_buddy_new(int level)
16 : {
17 4 : int size = 1 << level;
18 4 : struct dart_buddy * self =
19 4 : malloc(sizeof(struct dart_buddy) + sizeof(uint8_t) * (size * 2 - 2));
20 4 : self->level = level;
21 4 : memset(self->tree, NODE_UNUSED, size * 2 - 1);
22 4 : return self;
23 : }
24 :
25 : void
26 4 : dart_buddy_delete(struct dart_buddy * self) {
27 4 : free(self);
28 4 : }
29 :
30 : static inline int
31 24 : is_pow_of_2(uint32_t x) {
32 24 : return !(x & (x - 1));
33 : }
34 :
35 : static inline size_t
36 24 : next_pow_of_2(size_t x) {
37 24 : if (is_pow_of_2(x))
38 24 : return x;
39 0 : x |= x >> 1;
40 0 : x |= x >> 2;
41 0 : x |= x >> 4;
42 0 : x |= x >> 8;
43 0 : x |= x >> 16;
44 0 : x |= x >> 32;
45 0 : return x + 1;
46 : }
47 :
48 : static inline uint64_t
49 24 : _index_offset(int index, int level, int max_level) {
50 24 : return ((index + 1) - (1 << level)) << (max_level - level);
51 : }
52 :
53 : static void
54 24 : _mark_parent(struct dart_buddy * self, int index) {
55 : for (;;) {
56 24 : int buddy = index - 1 + (index & 1) * 2;
57 48 : if (buddy > 0 && (self->tree[buddy] == NODE_USED ||
58 24 : self->tree[buddy] == NODE_FULL)) {
59 0 : index = (index + 1) / 2 - 1;
60 0 : self->tree[index] = NODE_FULL;
61 : }
62 : else {
63 24 : return;
64 : }
65 0 : }
66 : }
67 :
68 : uint64_t
69 24 : dart_buddy_alloc(struct dart_buddy * self, size_t s) {
70 : int size;
71 24 : if (s == 0) {
72 0 : size = 1;
73 : }
74 : else {
75 24 : size = (int)next_pow_of_2(s);
76 : }
77 24 : int length = 1 << self->level;
78 :
79 24 : if (size > length)
80 0 : return -1;
81 :
82 24 : int index = 0;
83 24 : int level = 0;
84 :
85 575 : while (index >= 0) {
86 551 : if (size == length) {
87 24 : if (self->tree[index] == NODE_UNUSED) {
88 24 : self->tree[index] = NODE_USED;
89 24 : _mark_parent(self, index);
90 24 : return _index_offset(index, level, self->level);
91 : }
92 : }
93 : else {
94 : // size < length
95 527 : switch (self->tree[index]) {
96 : case NODE_USED:
97 : case NODE_FULL:
98 0 : break;
99 : case NODE_UNUSED:
100 : // split first
101 527 : self->tree[index] = NODE_SPLIT;
102 527 : self->tree[index * 2 + 1] = NODE_UNUSED;
103 527 : self->tree[index * 2 + 2] = NODE_UNUSED;
104 : default:
105 527 : index = index * 2 + 1;
106 527 : length /= 2;
107 527 : level++;
108 527 : continue;
109 : }
110 : }
111 0 : if (index & 1) {
112 0 : ++index;
113 0 : continue;
114 : }
115 : for (;;) {
116 0 : level--;
117 0 : length *= 2;
118 0 : index = (index + 1) / 2 - 1;
119 0 : if (index < 0)
120 0 : return -1;
121 0 : if (index & 1) {
122 0 : ++index;
123 0 : break;
124 : }
125 0 : }
126 : }
127 :
128 0 : return -1;
129 : }
130 :
131 : static void
132 551 : _combine(struct dart_buddy * self, int index) {
133 : for (;;) {
134 551 : int buddy = index - 1 + (index & 1) * 2;
135 551 : if (buddy < 0 || self->tree[buddy] != NODE_UNUSED) {
136 24 : self->tree[index] = NODE_UNUSED;
137 48 : while (((index = (index + 1) / 2 - 1) >= 0) &&
138 0 : self->tree[index] == NODE_FULL){
139 0 : self->tree[index] = NODE_SPLIT;
140 : }
141 48 : return;
142 : }
143 527 : index = (index + 1) / 2 - 1;
144 527 : }
145 : }
146 :
147 24 : int dart_buddy_free(struct dart_buddy * self, uint64_t offset)
148 : {
149 24 : int length = 1 << self->level;
150 24 : uint64_t left = 0;
151 24 : int index = 0;
152 :
153 24 : if (offset >= (uint64_t)length) {
154 0 : assert(offset < (uint64_t)length);
155 0 : return -1;
156 : }
157 :
158 : for (;;) {
159 551 : switch (self->tree[index]) {
160 : case NODE_USED:
161 24 : if (offset != left){
162 0 : assert (offset == left);
163 0 : return -1;
164 : }
165 24 : _combine(self, index);
166 24 : return 0;
167 : case NODE_UNUSED:
168 0 : assert (0);
169 : return -1;
170 : default:
171 527 : length /= 2;
172 527 : if (offset < left + length) {
173 527 : index = index * 2 + 1;
174 : }
175 : else {
176 0 : left += length;
177 0 : index = index * 2 + 2;
178 : }
179 527 : break;
180 : }
181 527 : }
182 : }
183 :
184 0 : int buddy_size(struct dart_buddy * self, uint64_t offset)
185 : {
186 0 : uint64_t left = 0;
187 0 : int length = 1 << self->level;
188 0 : int index = 0;
189 :
190 0 : assert(offset < (uint64_t)length);
191 :
192 : for (;;) {
193 0 : switch (self->tree[index]) {
194 : case NODE_USED:
195 0 : assert(offset == left);
196 0 : return length;
197 : case NODE_UNUSED:
198 0 : assert(0);
199 : return length;
200 : default:
201 0 : length /= 2;
202 0 : if (offset < left + length) {
203 0 : index = index * 2 + 1;
204 : }
205 : else {
206 0 : left += length;
207 0 : index = index * 2 + 2;
208 : }
209 0 : break;
210 : }
211 0 : }
212 : }
213 :
214 : static void
215 0 : _dump(struct dart_buddy * self, int index, int level) {
216 0 : switch (self->tree[index]) {
217 : case NODE_UNUSED:
218 0 : printf("(%"PRIu64":%d)",
219 : _index_offset(index, level, self->level),
220 0 : 1 << (self->level - level));
221 0 : break;
222 : case NODE_USED:
223 0 : printf("[%"PRIu64":%d]",
224 : _index_offset(index, level, self->level),
225 0 : 1 << (self->level - level));
226 0 : break;
227 : case NODE_FULL:
228 0 : printf("{");
229 0 : _dump(self, index * 2 + 1, level + 1);
230 0 : _dump(self, index * 2 + 2, level + 1);
231 0 : printf("}");
232 0 : break;
233 : default:
234 0 : printf("(");
235 0 : _dump(self, index * 2 + 1, level + 1);
236 0 : _dump(self, index * 2 + 2, level + 1);
237 0 : printf(")");
238 0 : break;
239 : }
240 0 : }
241 :
242 : void
243 0 : buddy_dump(struct dart_buddy * self) {
244 0 : _dump(self, 0, 0);
245 0 : printf("\n");
246 0 : }
|
