acadia/lib/glacier/container/binary_tree.h

#pragma once

#include "glacier/container/optional.h"
#include "glacier/container/pair.h"
#include "glacier/memory/move.h"
#include "glacier/memory/ref_counted.h"
#include "glacier/memory/ref_ptr.h"
#include "glacier/memory/reference.h"
#include "glacier/memory/unique_ptr.h"

namespace glcr {

template <typename K, typename V>
class BinaryTree {
 public:
  BinaryTree() = default;
  BinaryTree(const BinaryTree&) = delete;
  // FIXME: Implement move.
  BinaryTree(BinaryTree&&) = delete;

  void Insert(K key, V&& value);
  void Delete(K key);

  Optional<Ref<V>> Predecessor(K key);
  Optional<Ref<V>> Successor(K key);

  Optional<Ref<V>> Find(K key);

 private:
  // TODO: Consider adding a sharedptr type to
  // avoid making this "RefCounted".
  struct BinaryNode : public RefCounted<BinaryNode> {
    K key;
    V value;

    RefPtr<BinaryNode> left;
    RefPtr<BinaryNode> right;
    RefPtr<BinaryNode> parent;

    BinaryNode(K k, V v) : key(k), value(Move(v)) {}
  };

  RefPtr<BinaryNode> root_;

  // If this node exists, return it. Otherwise, this
  // will be the parent of where this node would be inserted.
  RefPtr<BinaryNode> FindOrInsertionParent(K key);
};

template <typename K, typename V>
void BinaryTree<K, V>::Insert(K key, V&& value) {
  auto parent = FindOrInsertionParent(key);
  if (parent.empty()) {
    root_ = AdoptPtr(new BinaryNode(key, Move(value)));
    return;
  }

  if (parent->key > key) {
    parent->left = AdoptPtr(new BinaryNode(key, Move(value)));
    parent->left->parent = parent;
  } else if (parent->key < key) {
    parent->right = AdoptPtr(new BinaryNode(key, Move(value)));
    parent->right->parent = parent;
  } else {
    parent->value = Move(value);
  }
}

template <typename K, typename V>
void BinaryTree<K, V>::Delete(K key) {
  auto node = FindOrInsertionParent(key);
  if (node.empty() || node->key != key) {
    return;
  }

  RefPtr<BinaryNode> new_child = nullptr;
  if (!node.left) {
    // No children.
    // Right child only.
    new_child = node.right;
  } else if (!node.right) {
    // Left child only.
    new_child = node.left;
  } else {
    // Find Successor.
    auto successor = node.right;
    while (successor.left) {
      successor = successor.left;
    }
    new_child = successor;
    if (successor != node.right) {
      successor.parent.left = successor.right;
    }
  }

  if (node == root_) {
    root_ = new_child;
  } else {
    if (node.parent.right == node) {
      node.parent.right = new_child;
    } else {
      node.parent.left = new_child;
    }
  }
}

template <typename K, typename V>
Optional<Ref<V>> BinaryTree<K, V>::Predecessor(K key) {
  auto current = FindOrInsertionParent(key);

  // The case where the current is the insertion parent and
  // the predecessor is unique. If the key was going to be
  // inserted as the left child, it shares its predecessor with the parent.
  if (current->key < key) {
    return Optional<Ref<V>>(current->value);
  }

  // Case where the predecessor is below us in the tree.
  if (current->left) {
    current = current->left;
    while (current->right) {
      current = current->right;
    }
    return {current->value};
  }

  // Case where the predecessor is above us in the tree.
  auto parent = current->parent;
  while (parent && (parent->left == current)) {
    current = parent;
    parent = current->parent;
  }
  if (parent) {
    return {parent->value};
  }
  return {};
}

template <typename K, typename V>
Optional<Ref<V>> BinaryTree<K, V>::Successor(K key) {
  auto current = FindOrInsertionParent(key);

  // The case where the current is the insertion parent and
  // the predecessor is unique. If the key was going to be
  // inserted as the left child, it shares its predecessor with the parent.
  if (current->key > key) {
    return Optional<Ref<V>>(current->value);
  }

  // Case where the predecessor is below us in the tree.
  if (current->right) {
    current = current->right;
    while (current->left) {
      current = current->left;
    }
    return {current->value};
  }

  // Case where the predecessor is above us in the tree.
  auto parent = current->parent;
  while (parent && (parent->right == current)) {
    current = parent;
    parent = current->parent;
  }
  if (parent) {
    return {parent->value};
  }
  return {};
}

template <typename K, typename V>
Optional<Ref<V>> BinaryTree<K, V>::Find(K key) {
  auto current = FindOrInsertionParent(key);
  if (current->key == key) {
    return Optional<Ref<V>>(current->value);
  }
  return {};
}

template <typename K, typename V>
RefPtr<typename BinaryTree<K, V>::BinaryNode>
BinaryTree<K, V>::FindOrInsertionParent(K key) {
  if (root_.empty()) {
    return nullptr;
  }

  auto current = root_;
  while (true) {
    if (key == current->key) {
      return current;
    } else if (key < current->key) {
      if (!current->left) {
        return current;
      }
      current = current->left;
    } else {
      if (!current->right) {
        return current;
      }
      current = current->right;
    }
  }
}

}  // namespace glcr