libvieter/src/tree/vieter_tree_balancing.c

#include "vieter_tree_balancing.h"
#include <string.h>

bool vieter_tree_node_validate(vieter_tree_node *node,
                               uint64_t passed_black_nodes,
                               uint64_t expected_black_nodes) {
  // Path to NULL children should contain the same amount of black nodes
  if (node == NULL) {
    return passed_black_nodes == expected_black_nodes;
  }

  if (vieter_tree_node_get(node, vieter_tree_node_black)) {
    passed_black_nodes++;
  }

  // Either the node itself is black, or its children are both either NULL or
  // black.
  bool correctly_colored_children =
      vieter_tree_node_get(node, vieter_tree_node_black) ||
      ((node->children[0] == NULL ||
        vieter_tree_node_get(node->children[0], vieter_tree_node_black)) &&
       (node->children[1] == NULL ||
        vieter_tree_node_get(node->children[1], vieter_tree_node_black)));

  return correctly_colored_children &&
         vieter_tree_node_validate(node->children[0], passed_black_nodes,
                                   expected_black_nodes) &&
         vieter_tree_node_validate(node->children[1], passed_black_nodes,
                                   expected_black_nodes);
}

vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
  vieter_tree_node *parent = node->parent;
  vieter_tree_node *grand_parent = parent->parent;
  vieter_tree_node *root = grand_parent;
  vieter_tree_node *left, *right;
  vieter_tree_node *children[4];

  uint64_t key_root = root->key;
  void *data_root = root->data;

  if (vieter_tree_node_get(node, vieter_tree_node_right)) {
    left = parent;
    right = node;

    if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
      root->key = parent->key;
      root->data = parent->data;

      parent->key = key_root;
      parent->data = data_root;

      children[0] = grand_parent->children[0];
      children[1] = parent->children[0];
      children[2] = node->children[0];
      children[3] = node->children[1];
    } else {
      root->key = node->key;
      root->data = node->data;

      node->key = key_root;
      node->data = data_root;

      children[0] = parent->children[0];
      children[1] = node->children[0];
      children[2] = node->children[1];
      children[3] = grand_parent->children[1];
    }
  } else {
    left = node;
    right = parent;

    if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
      root->key = node->key;
      root->data = node->data;

      node->key = key_root;
      node->data = data_root;

      children[0] = grand_parent->children[0];
      children[1] = node->children[0];
      children[2] = node->children[1];
      children[3] = parent->children[1];
    } else {
      root->key = parent->key;
      root->data = parent->data;

      parent->key = key_root;
      parent->data = data_root;

      children[0] = node->children[0];
      children[1] = node->children[1];
      children[2] = parent->children[1];
      children[3] = grand_parent->children[1];
    }
  }

  vieter_tree_node_set(
      root, vieter_tree_node_right,
      vieter_tree_node_get(grand_parent, vieter_tree_node_right));

  vieter_tree_node_set(left, vieter_tree_node_black, true);
  vieter_tree_node_set(left, vieter_tree_node_right, false);
  vieter_tree_node_set(right, vieter_tree_node_black, true);
  vieter_tree_node_set(left, vieter_tree_node_right, true);

  vieter_tree_node_set(root, vieter_tree_node_black, false);

  memcpy(&left->children, children, 2 * sizeof(vieter_tree_node *));

  if (children[0] != NULL) {
    children[0]->parent = left;
    vieter_tree_node_set(children[0], vieter_tree_node_right, false);
  }

  if (children[1] != NULL) {
    children[1]->parent = left;
    vieter_tree_node_set(children[0], vieter_tree_node_right, true);
  }

  memcpy(&right->children, children + 2, 2 * sizeof(vieter_tree_node *));

  if (children[2] != NULL) {
    children[2]->parent = right;
    vieter_tree_node_set(children[2], vieter_tree_node_right, false);
  }

  if (children[3] != NULL) {
    children[3]->parent = right;
    vieter_tree_node_set(children[2], vieter_tree_node_right, true);
  }

  root->children[0] = left;
  root->children[1] = right;
  left->parent = root;
  right->parent = root;

  return root;
}

void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
  while (!vieter_tree_node_get(node, vieter_tree_node_black) &&
         node->parent != NULL &&
         !vieter_tree_node_get(node->parent, vieter_tree_node_black)) {
    node = vieter_tree_node_balance(node);
  }

  // The root must always be black
  if (node->parent == NULL) {
    vieter_tree_node_set(node, vieter_tree_node_black, true);
  }
}

void vieter_tree_node_remove_balanced(vieter_tree_node *node) {}