#include #include #include "allocator_internal.h" #ifdef DEBUG #include void debug_print_tree(int indent, void *p) { TreeAlloc *node = (TreeAlloc*) p; if (node != NULL) { debug_print_tree(indent + 1, node->left); for (int ii = 0; ii < indent; ii++) { printf(" "); } if (node->color == COLOR_RED) { printf("\e[31"); } printf("%p %lu\n", node, node->size); if (node->color == COLOR_RED) { printf("\e[37"); } debug_print_tree(indent + 1, node->right); } } #endif TreeAlloc *insert_node_at(void *address, uintptr_t padding, uintptr_t align, uintptr_t size) { return NULL; } /* * Search for the node whose allocated region contains an address. */ TreeAlloc *search_by_address(TreeAlloc *root, void *address) { TreeAlloc *head = root; while (1) { if (head > (TreeAlloc*) address) { if (head->left == NULL) { return NULL; } else { head = head->left; } } else { if (head->right == NULL || head->right > (TreeAlloc*) address) { return head; } else { head = head->right; } } } } static uintptr_t effective_size(TreeAlloc *head, uintptr_t padding, uintptr_t align) { return head->size - (align_after(head + padding, align) - (void*) head); } /* * This is the most optimistic estimate of size that we can use which also preserves the ordering over * the tree. I had planned to use effective_size before I realized that it would break the tree * ordering. */ static uintptr_t pessimistic_size(TreeAlloc *head, uintptr_t padding, uintptr_t align) { return head->size - padding - align + 1; } TreeAlloc *search_by_size(TreeAlloc *root, uintptr_t padding, uintptr_t align, uintptr_t size) { TreeAlloc *head = root; while (1) { uintptr_t esize = pessimistic_size(head, padding, align); if (esize < size) { if (head->right == NULL) { return NULL; } else { head = head->right; } } else { if (head->left == NULL || pessimistic_size(head->left, padding, align) < size) { return head; } else { head = head->left; } } } } TreeAlloc *get_sibling(TreeAlloc *ta) { TreeAlloc *p = ta->parent; if (!p) return NULL; else if (p->left == ta) return p->right; else return p->left; } void rotate_left(TreeAlloc **root_ptr, TreeAlloc *ta) { TreeAlloc *parent, *tmp; parent = ta->parent; tmp = ta->right; if (!tmp) return; ta->right = tmp->left; tmp->left = ta; ta->parent = tmp; if (ta->right) ta->right->parent = ta; if (parent == NULL) { *root_ptr = tmp; } else { if (ta == parent->left) parent->left = tmp; else parent->right = tmp; } tmp->parent = parent; } void rotate_right(TreeAlloc **root_ptr, TreeAlloc *ta) { TreeAlloc *parent, *tmp; parent = ta->parent; tmp = ta->left; if (!tmp) return; ta->left = tmp->right; tmp->right = ta; ta->parent = tmp; if (ta->left) ta->left->parent = ta; if (parent == NULL) { *root_ptr = tmp; } else { if (ta == parent->left) parent->left = tmp; else parent->right = tmp; } tmp->parent = parent; } void repair_tree_after_insert(TreeAlloc **root_ptr, TreeAlloc *ta) { TreeAlloc *parent = ta->parent; if (parent == NULL) { ta->color = COLOR_BLACK; return; } TreeAlloc *grandparent = parent->parent; TreeAlloc *uncle = get_sibling(parent); if (parent->color == COLOR_RED) { if (uncle != NULL && uncle->color == COLOR_RED) { parent->color = COLOR_BLACK; uncle->color = COLOR_BLACK; grandparent->color = COLOR_RED; repair_tree_after_insert(root_ptr, grandparent); } else { if (ta == parent->left && parent == grandparent->left) { rotate_left(root_ptr, parent); ta = ta->left; } else { rotate_right(root_ptr, parent); ta = ta->right; } parent = ta->parent; grandparent = parent->parent; if (ta == parent->left) { rotate_right(root_ptr, grandparent); } else { rotate_left(root_ptr, grandparent); } parent->color = COLOR_BLACK; grandparent->color = COLOR_RED; } } } void replace_node(TreeAlloc **root_ptr, TreeAlloc *node, TreeAlloc *replace) { if (!node->parent) { *root_ptr = replace; } else { if (node == node->parent->left) node->parent->left = replace; else node->parent->right = replace; } if (replace) replace->parent = node->parent; } void repair_after_remove(TreeAlloc **root_ptr, TreeAlloc *node) { if (node->color == COLOR_RED) { node->color = COLOR_BLACK; } else { TreeAlloc *sibling = get_sibling(node); if (sibling->color == COLOR_RED) { if (node->parent->left == node) rotate_left(root_ptr, node->parent); else rotate_right(root_ptr, node->parent); node->parent->parent->color = node->parent->color = COLOR_BLACK; } if (sibling->left->color == COLOR_BLACK && sibling->right->color == COLOR_BLACK) { node->color = COLOR_BLACK; sibling->color = COLOR_RED; repair_after_remove(root_ptr, node->parent); } else { if (node->parent->left == node && sibling->right->color == COLOR_BLACK) { rotate_right(root_ptr, sibling); sibling = get_sibling(node); sibling->color = COLOR_RED; sibling->right->color = COLOR_RED; rotate_left(root_ptr, node->parent); node->color = get_sibling(node->parent)->color = COLOR_BLACK; } else if (node->parent->right == node && sibling->left->color == COLOR_BLACK) { rotate_left(root_ptr, sibling); sibling = get_sibling(node); sibling->color = COLOR_RED; sibling->left->color = COLOR_RED; rotate_right(root_ptr, sibling); node->color = get_sibling(node->parent)->color = COLOR_BLACK; } node->parent->color ^= node->parent->parent->color; node->parent->parent->color ^= node->parent->color; node->parent->color ^= node->parent->parent->color; } } } void remove_node(TreeAlloc **root_ptr, TreeAlloc *node) { char do_repair = 0; TreeAlloc *replace; TreeAlloc *parent = node->parent; if (!node->left) { replace = node->right; do_repair = node->color == COLOR_BLACK; replace_node(root_ptr, node, replace); } else if (!node->right) { replace = node->left; do_repair = node->color == COLOR_BLACK; replace_node(root_ptr, node, replace); } else { TreeAlloc *tmp = node->right; while (tmp->left) tmp = tmp->left; replace = tmp->right; do_repair = tmp->color == COLOR_BLACK; if (tmp != node->right) { replace_node(root_ptr, tmp, replace); tmp->right = node->right; node->right->parent = tmp; } replace_node(root_ptr, node, tmp); tmp->color = node->color; tmp->left = node->left; node->left->parent = tmp; } if (do_repair && replace) { repair_after_remove(root_ptr, replace); } } // Inserts a node into an empty tree. void insert_singleton(TreeAlloc **root_ptr, TreeAlloc *to_insert) { *root_ptr = to_insert; to_insert->parent = NULL; repair_tree_after_insert(root_ptr, to_insert); } void insert_right(TreeAlloc** root_ptr, TreeAlloc* to_insert, TreeAlloc* after) { if (after->right != NULL) { after->right->parent = to_insert; to_insert->right = after->right; } after->right = to_insert; to_insert->parent = after; repair_tree_after_insert(root_ptr, to_insert); } void insert_left(TreeAlloc** root_ptr, TreeAlloc* to_insert, TreeAlloc* before) { if (before->left != NULL) { before->left->parent = to_insert; to_insert->left = before->left; } before->left = to_insert; to_insert->parent = before; repair_tree_after_insert(root_ptr, to_insert); } int add_new_region(Arena *arena, uintptr_t size, uintptr_t padding, uintptr_t align) { uintptr_t realsize = size + align + alignof(WatermarkAlloc) + padding - 1; if (realsize < MIN_NEW_MEM_SIZE) { realsize = MIN_NEW_MEM_SIZE; } FreeSpace *reg = (FreeSpace*) arena->get_new_region(realsize); if (reg == NULL) { arena->error("can't allocate a new memory region!"); return 0; } FreeSpace *newreg = align_after(reg, alignof(WatermarkAlloc)); newreg->left = NULL; newreg->right = NULL; realsize -= (void*) newreg - (void*) reg; realsize -= realsize % alignof(WatermarkAlloc); newreg->size = realsize; if (arena->root_freespace == NULL) { insert_singleton((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) newreg); } else { FreeSpace *head = arena->root_freespace; while (head->right != NULL) { head = head->right; } insert_right((TreeAlloc**) arena->root_freespace, (TreeAlloc*) newreg, (TreeAlloc*) newreg); } return 1; } void unalloc(Arena *arena, void *addr) { #ifdef DEBUG printf("==== UNALLOCATING ====\n"); printf("=== FREESPACE TREE ===\n"); debug_print_tree(0, arena->root_freespace); printf("=== TREEALLOC TREE ===\n"); debug_print_tree(0, arena->root_treealloc); printf("=== END OF TREES ====\n"); #endif if (arena->root_treealloc == NULL) { arena->error("attempt to unallocate when there are no allocations!"); return; } // Find the node this address belongs to TreeAlloc *node = search_by_address(arena->root_treealloc, addr); if (node == NULL) { arena->error("attempt to free memory outside any allocations!"); return; } // Handle the watermark allocator in this region if (node->type == RT_WATERMARK) { // TODO: handle watermark deallocation return; } // Get rid of it remove_node(&arena->root_treealloc, node); // If there's free space on either side of it, merge it with the free space into a bigger chunk of // free space. uintptr_t size = node->size; FreeSpace *start = (FreeSpace*) node; if (node->before != NULL && node->before->type == RT_FREESPACE) { start = (FreeSpace*) node->before; size += node->before->size; remove_node((TreeAlloc**) &arena->root_freespace, node->before); } if (node->after != NULL && node->after->type == RT_FREESPACE) { size += node->after->size; remove_node((TreeAlloc**) &arena->root_freespace, node->after); } start->type = RT_FREESPACE; start->size = size; // And finally, insert the resulting free space. if (arena->root_freespace == NULL) { insert_singleton((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) start); } else { TreeAlloc *insert_point = search_by_size((TreeAlloc*) arena->root_freespace, 0, 1, size); if (insert_point == NULL) { TreeAlloc *head = (TreeAlloc*) arena->root_freespace; while (head->right != NULL) { head = head->right; } insert_right((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) start, head); } else { insert_left((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) start, insert_point); } } } void *alloc(Arena *arena, uintptr_t size, uintptr_t align) { uintptr_t actual_align = lcm(alignof(struct WatermarkAlloc), align); #ifdef DEBUG printf("==== ALLOCATING =====\n"); printf("=== FREESPACE TREE ===\n"); debug_print_tree(0, arena->root_freespace); printf("=== TREEALLOC TREE ===\n"); debug_print_tree(0, arena->root_treealloc); printf("=== END OF TREES ====\n"); #endif if (arena->root_freespace == NULL) { // Handle being out of freespace. if (!add_new_region(arena, size, sizeof(TreeAlloc), actual_align)) { return NULL; } return alloc(arena, size, align); } else { TreeAlloc *region = search_by_size((TreeAlloc*) arena->root_freespace, sizeof(TreeAlloc), actual_align, size); if (region == NULL) { // Handle insufficient freespace or fragmentation. if (!add_new_region(arena, size, sizeof(TreeAlloc), actual_align)) { return NULL; } return alloc(arena, size, align); } remove_node((TreeAlloc**) &arena->root_freespace, region); void *true_end = align_after(align_after(((void*) region) + sizeof(TreeAlloc), actual_align) + size, alignof(WatermarkAlloc)); // The size of the new allocation (adjusted for region header and alignment uintptr_t new_size = true_end - (void*) region; // The size of the free space region following the new allocation uintptr_t new_free_size = region->size - new_size; region->right = NULL; region->left = NULL; region->type = RT_TREE_NODE; #ifdef DEBUG printf("sizeof(TreeAlloc): %lu\n", (uintptr_t) sizeof(TreeAlloc)); printf("start: %p, end: %p, adjusted end: %p\n", region, ((void*) region) + size, true_end); printf("size: %lu -> %lu\n", size, new_size); #endif if (arena->root_treealloc == NULL) { insert_singleton((TreeAlloc**) &arena->root_treealloc, region); } else { TreeAlloc *insert_point = search_by_address((TreeAlloc*) arena->root_treealloc, region); insert_right(&arena->root_treealloc, region, insert_point); } if (new_free_size >= sizeof(FreeSpace)) { // If there's enough free space after the allocation, use it! region->size = new_size; // Safe because the allocated region tree is not sorted by size. FreeSpace *new_free = (FreeSpace*) ((void*) region + new_size); new_free->left = NULL; new_free->right = NULL; new_free->type = RT_FREESPACE; new_free->size = new_free_size; if (arena->root_freespace == NULL) { insert_singleton((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) new_free); } else { FreeSpace *insert_point = (FreeSpace*) search_by_size((TreeAlloc*) arena->root_freespace, 0, 1, new_free_size); insert_left((TreeAlloc**) &arena->root_freespace, (TreeAlloc*) new_free, (TreeAlloc*) insert_point); } } return align_after(region + sizeof(TreeAlloc), actual_align); } } void *alloc_growable(Arena *arena, uintptr_t size, uintptr_t align) { // TODO: Basically the same as above, but put the allocated region in the center of the largest free // space. Due to alignment and whatnot, the code will be gory. return NULL; }