Merge branch 'tree_alloc' of gitea:cameron/cs241-project into tree_alloc

tree_alloc.c

@@ -160,30 +160,91 @@ void repair_tree_after_insert(TreeAlloc **root_ptr, TreeAlloc *ta) {
 	}
 }
 
-// TODO: Rewrite for self-balancing tree.
-void remove_node(TreeAlloc** root_ptr, TreeAlloc* node) {
-  TreeAlloc *replace = NULL;
-  if (node->left == NULL) {
-    replace = node->right;
-  } else if (node->right != NULL) {
-    replace = node->right;
-    TreeAlloc *head = node->left;
-    while (head->right != NULL) {
-      head = head->right;
-    }
-    head->right = head->parent->right;
-    head->right->parent = head;
-  }
-  if (node->parent == NULL) {
-    replace->parent = NULL;
-    *root_ptr = replace;
-  } else {
-    if (node == node->parent->left) {
-      node->parent->left = replace;
-    } else {
-      node->parent->right = replace;
-    }
-  }
+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, 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, 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_fix = node->color == COLOR_BLACK;
+		replace_node(root_ptr, node, replace);
+	} else if (!node->right) {
+		replace = node->left;
+		do_fix = 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_fix = tmp == 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) {
+		repair_after_remove(root_ptr, replace);
+	}
 }