You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

578 lines
17 KiB

2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
2 years ago
  1. #include <stddef.h>
  2. #include <stdalign.h>
  3. #include "allocator_internal.h"
  4. #ifdef DEBUG
  5. #include <stdio.h>
  6. void debug_print_tree(int indent, void *p) {
  7. TreeAlloc *node = (TreeAlloc*) p;
  8. if (node != NULL) {
  9. debug_print_tree(indent + 1, node->left);
  10. for (int ii = 0; ii < indent; ii++) { printf(" "); }
  11. if (node->color == COLOR_RED) { printf("\e[31m"); }
  12. printf("%p %lu\n", node, node->size);
  13. if (node->color == COLOR_RED) { printf("\e[37m"); }
  14. debug_print_tree(indent + 1, node->right);
  15. }
  16. }
  17. #endif
  18. TreeAlloc *insert_node_at(void *address, uintptr_t padding, uintptr_t align, uintptr_t size) {
  19. return NULL;
  20. }
  21. /*
  22. * Search for the node whose allocated region contains an address.
  23. */
  24. TreeAlloc *search_by_address(TreeAlloc *root, void *address) {
  25. TreeAlloc *head = root;
  26. while (1) {
  27. if (head > (TreeAlloc*) address) {
  28. if (head->left == NULL) {
  29. return NULL;
  30. } else {
  31. head = head->left;
  32. }
  33. } else {
  34. if (head->right == NULL || head->right > (TreeAlloc*) address) {
  35. return head;
  36. } else {
  37. head = head->right;
  38. }
  39. }
  40. }
  41. }
  42. static uintptr_t effective_size(TreeAlloc *head, uintptr_t padding, uintptr_t align) {
  43. return head->size - (align_after(head + padding, align) - (void*) head);
  44. }
  45. /*
  46. * This is the most optimistic estimate of size that we can use which also preserves the ordering over
  47. * the tree. I had planned to use effective_size before I realized that it would break the tree
  48. * ordering.
  49. */
  50. static uintptr_t pessimistic_size(TreeAlloc *head, uintptr_t padding, uintptr_t align) {
  51. return head->size - padding - align + 1;
  52. }
  53. TreeAlloc *search_by_size(TreeAlloc *root, uintptr_t padding, uintptr_t align, uintptr_t size) {
  54. TreeAlloc *head = root;
  55. while (1) {
  56. uintptr_t esize = pessimistic_size(head, padding, align);
  57. if (esize < size) {
  58. if (head->right == NULL) {
  59. return NULL;
  60. } else {
  61. head = head->right;
  62. }
  63. } else {
  64. if (head->left == NULL || pessimistic_size(head->left, padding, align) < size) {
  65. return head;
  66. } else {
  67. head = head->left;
  68. }
  69. }
  70. }
  71. }
  72. TreeAlloc *succ(TreeAlloc *el) {
  73. if (el->right != NULL) {
  74. el = el->right;
  75. while (el->left != NULL) {
  76. el = el->left;
  77. }
  78. return el;
  79. }
  80. while (el->parent != NULL && el == el->parent->right) {
  81. el = el->parent;
  82. }
  83. return el->parent;
  84. }
  85. TreeAlloc *pred(TreeAlloc *el) {
  86. if (el->left != NULL) {
  87. el = el->left;
  88. while (el->right != NULL) {
  89. el = el->right;
  90. }
  91. return el;
  92. }
  93. while (el->parent != NULL && el == el->parent->left) {
  94. el = el->parent;
  95. }
  96. return el->parent;
  97. }
  98. TreeAlloc *get_sibling(TreeAlloc *ta) {
  99. TreeAlloc *p = ta->parent;
  100. if (!p)
  101. return NULL;
  102. else if (p->left == ta)
  103. return p->right;
  104. else
  105. return p->left;
  106. }
  107. void rotate_left(TreeAlloc **root_ptr, TreeAlloc *ta) {
  108. TreeAlloc *parent, *tmp;
  109. parent = ta->parent;
  110. tmp = ta->right;
  111. if (!tmp) return;
  112. ta->right = tmp->left;
  113. tmp->left = ta;
  114. ta->parent = tmp;
  115. if (ta->right) ta->right->parent = ta;
  116. if (parent == NULL) {
  117. *root_ptr = tmp;
  118. } else {
  119. if (ta == parent->left)
  120. parent->left = tmp;
  121. else
  122. parent->right = tmp;
  123. }
  124. tmp->parent = parent;
  125. }
  126. void rotate_right(TreeAlloc **root_ptr, TreeAlloc *ta) {
  127. TreeAlloc *parent, *tmp;
  128. parent = ta->parent;
  129. tmp = ta->left;
  130. if (!tmp) return;
  131. ta->left = tmp->right;
  132. tmp->right = ta;
  133. ta->parent = tmp;
  134. if (ta->left) ta->left->parent = ta;
  135. if (parent == NULL) {
  136. *root_ptr = tmp;
  137. } else {
  138. if (ta == parent->left)
  139. parent->left = tmp;
  140. else
  141. parent->right = tmp;
  142. }
  143. tmp->parent = parent;
  144. }
  145. #define IS_BLACK_NODE(n) (n == NULL || n->color == COLOR_BLACK)
  146. #define IS_RED_NODE(n) (n != NULL && n->color == COLOR_RED)
  147. void repair_tree_after_insert(TreeAlloc **root_ptr, TreeAlloc *ta) {
  148. TreeAlloc *parent = ta->parent;
  149. #ifdef DEBUG
  150. printf("=== PRE-INSERT-FIXUP ===\n");
  151. printf("===== CURRENT TREE =====\n");
  152. debug_print_tree(0, *root_ptr);
  153. printf("===== END OF TREES =====\n");
  154. #endif
  155. if (parent == NULL) {
  156. ta->color = COLOR_BLACK;
  157. return;
  158. }
  159. TreeAlloc *grandparent = parent->parent;
  160. TreeAlloc *uncle = get_sibling(parent);
  161. if (IS_RED_NODE(parent)) {
  162. if (IS_RED_NODE(uncle)) {
  163. parent->color = COLOR_BLACK;
  164. uncle->color = COLOR_BLACK;
  165. grandparent->color = COLOR_RED;
  166. repair_tree_after_insert(root_ptr, grandparent);
  167. } else {
  168. if (ta == parent->left && parent == grandparent->left) {
  169. rotate_left(root_ptr, parent);
  170. ta = ta->left;
  171. parent = parent->left;
  172. } else {
  173. rotate_right(root_ptr, parent);
  174. ta = ta->right;
  175. parent = parent->right;
  176. }
  177. grandparent = parent->parent;
  178. if (ta == parent->left) {
  179. rotate_right(root_ptr, grandparent);
  180. } else {
  181. rotate_left(root_ptr, grandparent);
  182. }
  183. parent->color = COLOR_BLACK;
  184. grandparent->color = COLOR_RED;
  185. }
  186. }
  187. #ifdef DEBUG
  188. printf("== POST-INSERT-FIXUP ===\n");
  189. printf("===== CURRENT TREE =====\n");
  190. debug_print_tree(0, *root_ptr);
  191. printf("===== END OF TREES =====\n");
  192. #endif
  193. }
  194. void replace_node(TreeAlloc **root_ptr, TreeAlloc *node, TreeAlloc *replace) {
  195. if (!node->parent) {
  196. *root_ptr = replace;
  197. } else {
  198. if (node == node->parent->left)
  199. node->parent->left = replace;
  200. else
  201. node->parent->right = replace;
  202. }
  203. if (replace) replace->parent = node->parent;
  204. }
  205. void repair_after_remove(TreeAlloc **root_ptr, TreeAlloc *node) {
  206. if (IS_RED_NODE(node)) {
  207. node->color = COLOR_BLACK;
  208. } else {
  209. TreeAlloc *sibling = get_sibling(node);
  210. if (IS_RED_NODE(sibling)) {
  211. if (node->parent->left == node) {
  212. rotate_left(root_ptr, node->parent);
  213. } else {
  214. rotate_right(root_ptr, node->parent);
  215. }
  216. node->parent->parent->color = node->parent->color = COLOR_BLACK;
  217. }
  218. if (IS_BLACK_NODE(sibling->left) && IS_BLACK_NODE(sibling->right)) {
  219. node->color = COLOR_BLACK;
  220. sibling->color = COLOR_RED;
  221. repair_after_remove(root_ptr, node->parent);
  222. } else {
  223. if (node->parent->left == node && IS_BLACK_NODE(sibling)) {
  224. rotate_right(root_ptr, sibling);
  225. sibling = get_sibling(node);
  226. sibling->color = COLOR_RED;
  227. sibling->right->color = COLOR_RED;
  228. rotate_left(root_ptr, node->parent);
  229. node->color = get_sibling(node->parent)->color = COLOR_BLACK;
  230. } else if (node->parent->right == node && IS_BLACK_NODE(sibling)) {
  231. rotate_left(root_ptr, sibling);
  232. sibling = get_sibling(node);
  233. sibling->color = COLOR_RED;
  234. sibling->left->color = COLOR_RED;
  235. rotate_right(root_ptr, sibling);
  236. node->color = get_sibling(node->parent)->color = COLOR_BLACK;
  237. }
  238. node->parent->color ^= node->parent->parent->color;
  239. node->parent->parent->color ^= node->parent->color;
  240. node->parent->color ^= node->parent->parent->color;
  241. }
  242. }
  243. }
  244. void remove_node(TreeAlloc **root_ptr, TreeAlloc *node) {
  245. char do_repair = 0;
  246. #ifdef DEBUG
  247. printf("====== PRE-REMOVE ======\n");
  248. printf("===== CURRENT TREE =====\n");
  249. debug_print_tree(0, *root_ptr);
  250. printf("===== END OF TREES =====\n");
  251. #endif
  252. TreeAlloc *replace;
  253. TreeAlloc *parent = node->parent;
  254. if (!node->left) {
  255. replace = node->right;
  256. do_repair = node->color == COLOR_BLACK;
  257. replace_node(root_ptr, node, replace);
  258. } else if (!node->right) {
  259. replace = node->left;
  260. do_repair = node->color == COLOR_BLACK;
  261. replace_node(root_ptr, node, replace);
  262. } else {
  263. TreeAlloc *tmp = node->right;
  264. while (tmp->left) tmp = tmp->left;
  265. replace = tmp->right;
  266. do_repair = tmp->color == COLOR_BLACK;
  267. if (tmp != node->right) {
  268. replace_node(root_ptr, tmp, replace);
  269. tmp->right = node->right;
  270. node->right->parent = tmp;
  271. }
  272. replace_node(root_ptr, node, tmp);
  273. tmp->color = node->color;
  274. tmp->left = node->left;
  275. node->left->parent = tmp;
  276. }
  277. if (do_repair && replace && replace->parent != NULL) {
  278. #ifdef DEBUG
  279. printf("=== PRE-REMOVE-FIXUP ===\n");
  280. printf("===== CURRENT TREE =====\n");
  281. debug_print_tree(0, *root_ptr);
  282. printf("===== END OF TREES =====\n");
  283. #endif
  284. repair_after_remove(root_ptr, replace);
  285. }
  286. }
  287. // Inserts a node into an empty tree.
  288. void insert_singleton(TreeAlloc **root_ptr, TreeAlloc *to_insert) {
  289. #ifdef DEBUG
  290. printf("= PRE-INSERT-SINGLETON =\n");
  291. printf("===== CURRENT TREE =====\n");
  292. debug_print_tree(0, *root_ptr);
  293. printf("===== END OF TREES =====\n");
  294. #endif
  295. *root_ptr = to_insert;
  296. to_insert->parent = NULL;
  297. repair_tree_after_insert(root_ptr, to_insert);
  298. }
  299. void insert_right(TreeAlloc** root_ptr, TreeAlloc* to_insert, TreeAlloc* after) {
  300. #ifdef DEBUG
  301. printf("=== PRE-INSERT-RIGHT ===\n");
  302. printf("===== CURRENT TREE =====\n");
  303. debug_print_tree(0, *root_ptr);
  304. printf("===== END OF TREES =====\n");
  305. #endif
  306. if (after->right != NULL) {
  307. after->right->parent = to_insert;
  308. to_insert->right = after->right;
  309. }
  310. after->right = to_insert;
  311. to_insert->parent = after;
  312. repair_tree_after_insert(root_ptr, to_insert);
  313. }
  314. void insert_left(TreeAlloc** root_ptr, TreeAlloc* to_insert, TreeAlloc* before) {
  315. #ifdef DEBUG
  316. printf("=== PRE-INSERT-LEFT ====\n");
  317. printf("===== CURRENT TREE =====\n");
  318. debug_print_tree(0, *root_ptr);
  319. printf("===== END OF TREES =====\n");
  320. #endif
  321. if (before->left != NULL) {
  322. before->left->parent = to_insert;
  323. to_insert->left = before->left;
  324. }
  325. before->left = to_insert;
  326. to_insert->parent = before;
  327. repair_tree_after_insert(root_ptr, to_insert);
  328. }
  329. int add_new_region(Arena *arena, uintptr_t size, uintptr_t padding, uintptr_t align) {
  330. uintptr_t realsize = size + align + alignof(WatermarkAlloc) + padding - 1;
  331. if (realsize < MIN_NEW_MEM_SIZE) {
  332. realsize = MIN_NEW_MEM_SIZE;
  333. }
  334. FreeSpace *reg = (FreeSpace*) arena->get_new_region(realsize);
  335. if (reg == NULL) {
  336. arena->error("can't allocate a new memory region!");
  337. return 0;
  338. }
  339. FreeSpace *newreg = align_after(reg, alignof(WatermarkAlloc));
  340. newreg->left = NULL;
  341. newreg->right = NULL;
  342. realsize -= (void*) newreg - (void*) reg;
  343. realsize -= realsize % alignof(WatermarkAlloc);
  344. newreg->size = realsize;
  345. if (arena->root_freespace == NULL) {
  346. insert_singleton((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) newreg);
  347. } else {
  348. FreeSpace *head = arena->root_freespace;
  349. while (head->right != NULL) {
  350. head = head->right;
  351. }
  352. insert_right((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) newreg, (TreeAlloc*) head);
  353. }
  354. #ifdef DEBUG
  355. printf("= POST-REGION-CREATION =\n");
  356. printf("==== FREESPACE TREE ====\n");
  357. debug_print_tree(0, arena->root_freespace);
  358. printf("==== TREEALLOC TREE ====\n");
  359. debug_print_tree(0, arena->root_treealloc);
  360. printf("===== END OF TREES =====\n");
  361. #endif
  362. return 1;
  363. }
  364. void unalloc(Arena *arena, void *addr) {
  365. #ifdef DEBUG
  366. printf("==== UNALLOCATING ====\n");
  367. printf("=== FREESPACE TREE ===\n");
  368. debug_print_tree(0, arena->root_freespace);
  369. printf("=== TREEALLOC TREE ===\n");
  370. debug_print_tree(0, arena->root_treealloc);
  371. printf("==== END OF TREES ====\n");
  372. #endif
  373. if (arena->root_treealloc == NULL) {
  374. arena->error("attempt to unallocate when there are no allocations!");
  375. return;
  376. }
  377. // Find the node this address belongs to
  378. TreeAlloc *node = search_by_address(arena->root_treealloc, addr);
  379. if (node == NULL) {
  380. arena->error("attempt to free memory outside any allocations!");
  381. return;
  382. }
  383. // Handle the watermark allocator in this region
  384. if (node->type == RT_WATERMARK) {
  385. // TODO: handle watermark deallocation
  386. return;
  387. }
  388. // Get rid of it
  389. remove_node(&arena->root_treealloc, node);
  390. // If there's free space on either side of it, merge it with the free space into a bigger chunk of
  391. // free space.
  392. uintptr_t size = node->size;
  393. FreeSpace *start = (FreeSpace*) node;
  394. if (node->before != NULL && node->before->type == RT_FREESPACE) {
  395. start = (FreeSpace*) node->before;
  396. size += node->before->size;
  397. remove_node((TreeAlloc**) &arena->root_freespace, node->before);
  398. }
  399. if (node->after != NULL && node->after->type == RT_FREESPACE) {
  400. size += node->after->size;
  401. remove_node((TreeAlloc**) &arena->root_freespace, node->after);
  402. }
  403. start->type = RT_FREESPACE;
  404. start->size = size;
  405. // And finally, insert the resulting free space.
  406. if (arena->root_freespace == NULL) {
  407. insert_singleton((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) start);
  408. } else {
  409. TreeAlloc *insert_point = search_by_size((TreeAlloc*) arena->root_freespace, 0, 1, size);
  410. if (insert_point == NULL) {
  411. TreeAlloc *head = (TreeAlloc*) arena->root_freespace;
  412. while (head->right != NULL) {
  413. head = head->right;
  414. }
  415. insert_right((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) start, head);
  416. } else {
  417. insert_left((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) start, insert_point);
  418. }
  419. }
  420. }
  421. void *alloc(Arena *arena, uintptr_t size, uintptr_t align) {
  422. uintptr_t actual_align = lcm(alignof(struct WatermarkAlloc), align);
  423. #ifdef DEBUG
  424. printf("==== ALLOCATING =====\n");
  425. printf("=== FREESPACE TREE ===\n");
  426. debug_print_tree(0, arena->root_freespace);
  427. printf("=== TREEALLOC TREE ===\n");
  428. debug_print_tree(0, arena->root_treealloc);
  429. printf("==== END OF TREES ====\n");
  430. #endif
  431. if (arena->root_freespace == NULL) {
  432. // Handle being out of freespace.
  433. #ifdef DEBUG
  434. printf("Out of freespace nodes; getting more\n");
  435. #endif
  436. if (!add_new_region(arena, size, sizeof(TreeAlloc), actual_align)) {
  437. return NULL;
  438. }
  439. return alloc(arena, size, align);
  440. } else {
  441. TreeAlloc *region = search_by_size((TreeAlloc*) arena->root_freespace, sizeof(TreeAlloc), actual_align, size);
  442. if (region == NULL) {
  443. // Handle insufficient freespace or fragmentation.
  444. #ifdef DEBUG
  445. printf("Out of sufficiently large freespace nodes; getting more\n");
  446. #endif
  447. if (!add_new_region(arena, size, sizeof(TreeAlloc), actual_align)) {
  448. return NULL;
  449. }
  450. return alloc(arena, size, align);
  451. }
  452. remove_node((TreeAlloc**) &arena->root_freespace, region);
  453. void *true_end = align_after(align_after(((void*) region) + sizeof(TreeAlloc), actual_align) + size, alignof(WatermarkAlloc));
  454. // The size of the new allocation (adjusted for region header and alignment
  455. uintptr_t new_size = true_end - (void*) region;
  456. // The size of the free space region following the new allocation
  457. uintptr_t new_free_size = region->size - new_size;
  458. region->right = NULL;
  459. region->left = NULL;
  460. region->type = RT_TREE_NODE;
  461. #ifdef DEBUG
  462. printf("start: %p, end: %p, adjusted end: %p\n", region, ((void*) region) + size, true_end);
  463. printf("size: %lu -> %lu\n", size, new_size);
  464. #endif
  465. if (arena->root_treealloc == NULL) {
  466. insert_singleton((TreeAlloc**) &arena->root_treealloc, region);
  467. } else {
  468. #ifdef DEBUG
  469. printf("searching for an insert point\n");
  470. #endif
  471. TreeAlloc *insert_point = search_by_address((TreeAlloc*) arena->root_treealloc, region);
  472. if (insert_point == NULL) {
  473. TreeAlloc *head = arena->root_treealloc;
  474. while (head->left != NULL) {
  475. head = head->left;
  476. }
  477. #ifdef DEBUG
  478. printf("none found; inserting before %p\n", head);
  479. #endif
  480. insert_left(&arena->root_treealloc, region, head);
  481. } else {
  482. #ifdef DEBUG
  483. printf("found one: %p\n", insert_point);
  484. #endif
  485. insert_right(&arena->root_treealloc, region, insert_point);
  486. }
  487. }
  488. if (new_free_size >= sizeof(FreeSpace)) {
  489. // If there's enough free space after the allocation, use it!
  490. region->size = new_size; // Safe because the allocated region tree is not sorted by size.
  491. FreeSpace *new_free = (FreeSpace*) ((void*) region + new_size);
  492. new_free->left = NULL;
  493. new_free->right = NULL;
  494. new_free->type = RT_FREESPACE;
  495. new_free->size = new_free_size;
  496. if (arena->root_freespace == NULL) {
  497. insert_singleton((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) new_free);
  498. } else {
  499. FreeSpace *insert_point = (FreeSpace*) search_by_size((TreeAlloc*) arena->root_freespace, 0, 1, new_free_size);
  500. insert_left((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) new_free, (TreeAlloc*) insert_point);
  501. }
  502. // Set the region following this one to be the new free space
  503. region->after = (TreeAlloc*) new_free;
  504. } else {
  505. // There isn't a free space after this one, so put the `next` pointer at the next allocated
  506. // region.
  507. region->after = succ(region);
  508. }
  509. // I seem to have forgotten about the fact that memory may not be contiguous
  510. if (region->after != NULL && region->after != (void*) region + region->size) {
  511. region->after = NULL;
  512. }
  513. // Also make sure the `before` pointer is correct.
  514. TreeAlloc *before_alloc = pred(region);
  515. if (before_alloc == NULL || ((void*) before_alloc) + before_alloc->size < (void*) region) {
  516. region->before = search_by_address((TreeAlloc*) arena->root_freespace, region);
  517. } else {
  518. region->before = before_alloc;
  519. }
  520. // I seem to have forgotten about the fact that memory may not be contiguous
  521. if (region->before != NULL && region->before != (void*) region->before + region->before->size) {
  522. region->before = NULL;
  523. }
  524. #ifdef DEBUG
  525. printf("region is still at %p\n", region);
  526. #endif
  527. return align_after((void*) region + sizeof(TreeAlloc), actual_align);
  528. }
  529. }
  530. void *alloc_growable(Arena *arena, uintptr_t size, uintptr_t align) {
  531. // TODO: Basically the same as above, but put the allocated region in the center of the largest free
  532. // space. Due to alignment and whatnot, the code will be gory.
  533. return NULL;
  534. }