libvieter/src/tree/vieter_tree_node.c

#include "vieter_tree_node.h"
#include "vieter_tree_balancing.h"

#include <string.h>

vieter_tree_node *vieter_tree_node_init() {
  return calloc(1, sizeof(vieter_tree_node));
}

void vieter_tree_node_free(vieter_tree_node *node) { free(node); }

void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
                                vieter_tree_node *child) {
  if (parent == NULL) {
    return;
  }

  bool side = key > parent->key;
  parent->children[side] = child;

  if (child != NULL) {
    child->parent = parent;
    vieter_tree_node_set(child, vieter_tree_node_right, side);
  }
}

void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace,
                                         vieter_tree_node *replacement) {
  to_replace->key = replacement->key;
  to_replace->data = replacement->data;
  memcpy(to_replace->children, replacement->children,
         2 * sizeof(vieter_tree_node *));

  if (to_replace->children[0] != NULL)
    to_replace->children[0]->parent = to_replace;

  if (to_replace->children[1] != NULL)
    to_replace->children[1]->parent = to_replace;
}

vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
                                          void *data) {
  vieter_tree_node *node = root;
  vieter_tree_node *parent = root;

  while (node != NULL) {
    if (node->key == key) {
      return vieter_tree_already_present;
    }

    parent = node;
    node = node->children[key > parent->key];
  }

  vieter_tree_node *new_node = vieter_tree_node_init();
  new_node->key = key;
  new_node->data = data;

  vieter_tree_node_add_child(parent, key, new_node);
  vieter_tree_node_balance_after_insert(new_node);

  return vieter_tree_ok;
}

vieter_tree_error vieter_tree_node_search_node(vieter_tree_node **out,
                                               vieter_tree_node *root,
                                               uint64_t key) {
  vieter_tree_node *node = root;

  while (node != NULL) {
    if (node->key == key) {
      *out = node;

      return vieter_tree_ok;
    }

    node = node->children[key > node->key];
  }

  return vieter_tree_not_present;
}

vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root,
                                          uint64_t key) {
  vieter_tree_node *target;
  vieter_tree_error res = vieter_tree_node_search_node(&target, root, key);

  if (res != vieter_tree_ok) {
    return res;
  }

  *out = target->data;

  return vieter_tree_ok;
}

vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
                                          uint64_t key) {
  vieter_tree_node *target;
  vieter_tree_error res = vieter_tree_node_search_node(&target, root, key);

  if (res != vieter_tree_ok) {
    return res;
  }

  *out = target->data;

  if (target->children[0] == NULL && target->children[1] == NULL) {
    vieter_tree_node_add_child(target->parent, target->key, NULL);

    vieter_tree_node_free(target);
  } else if ((target->children[0] == NULL) ^ (target->children[1] == NULL)) {
    vieter_tree_node *child =
        target->children[0] != NULL ? target->children[0] : target->children[1];

    if (target->parent != NULL) {
      vieter_tree_node_add_child(target->parent, child->key, child);
      vieter_tree_node_free(target);
    } else {
      vieter_tree_node_replace_with_child(target, child);
      vieter_tree_node_free(child);
    }

  } else {
    vieter_tree_node *replacement = target->children[1];

    while (replacement->children[0] != NULL) {
      replacement = replacement->children[0];
    }

    // We use replacement->key here because the right child can be NULL, so
    // reading its key isn't safe. Using replacement->key however, the child
    // will still get placed into the right location because of how binary
    // trees work.
    vieter_tree_node_add_child(replacement->parent, replacement->key,
                               replacement->children[1]);

    target->key = replacement->key;
    target->data = replacement->data;

    vieter_tree_node_free(replacement);
  }

  return vieter_tree_ok;
}

void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag,
                          bool set) {
  if (set) {
    node->flags |= flag;
  } else {
    node->flags &= ~flag;
  }
}

bool vieter_tree_node_get(vieter_tree_node *node, vieter_tree_node_flag flag) {
  return (node->flags & flag) != 0;
}