cs440-assignment2/Map.hpp

#ifndef _POWELL_CS440
#define _POWELL_CS440
#include <algorithm>
#include <cassert>
#include <iostream>
#include <memory>
#include <optional>
#include <utility>
namespace cs440 {
// universal type defs here
namespace {
enum class Direction { Left, Right };
Direction operator!(Direction dir) {
  switch (dir) {
  case Direction::Left:
    return Direction::Right;
  case Direction::Right:
    return Direction::Left;
  default:
    assert(false);
  }
}
int x5 = 5;
int x6 = 6;
int x7 = 7;
enum class Color { Red, Black };
} // namespace
template <typename Key_T, typename Mapped_T> class Map {
  // Type definitions here
  using ValueType = std::pair<const Key_T, Mapped_T>;
  using internal_ValueType = std::pair<Key_T, Mapped_T>;

  struct Node {
    Node *parent;
    internal_ValueType val;
    std::unique_ptr<Node> left;
    std::unique_ptr<Node> right;
    Color color;
    Node *prev;
    Node *next;
    Map *map;
    Node(internal_ValueType val, Map *map)
        : parent{nullptr}, val{val}, left{}, right{}, color{Color::Red},
          prev{nullptr}, next{nullptr}, map{map} {}
    Node(const Node &rhs)
        : parent{nullptr}, val{rhs.val},
          left{std::make_unique<Node>(*rhs.left)},
          right{std::make_unique<Node>(*rhs.right)}, color{rhs.color},
          prev{nullptr}, next{nullptr}, map{rhs.map} {
      if (this->left) {
        this->left->parent = this;
      }
      if (this->right) {
        this->right->parent = this;
      }
    }
    Node(Node &&rhs)
        : parent{nullptr}, val{std::move(rhs.val)}, left{std::move(rhs.left)},
          right{std::move(rhs.right)}, color{rhs.color}, prev{nullptr},
          next{nullptr}, map{rhs.map} {
      if (this->left) {
        this->left->parent = this;
      }
      if (this->right) {
        this->right->parent = this;
      }
    }
    ~Node() {}
    Node &operator=(const Node &rhs) {
      // retain parent as is, common case is the copy or move is happening due
      // to a rotation where parent can get wonky
      // this->parent
      this->val = rhs.val;
      this->left = std::make_unique<Node>(*rhs.left);
      this->right = std::make_unique<Node>(*rhs.right);
      this->color = rhs.color;
      if (this->left) {
        this->left->parent = this;
        this->left->restore_ordering();
      }
      if (this->right) {
        this->right->parent = this;
        this->right->restore_ordering();
      }

      this->map = rhs.map;
      return *this;
    }
    Node &operator=(Node &&rhs) {
      // retain parent as is, common case is the copy or move is happening due
      // to a rotation where parent can get wonky
      // this->parent
      this->val = rhs.val;
      this->left = std::move(rhs.left);
      this->right = std::move(rhs.right);
      this->color = rhs.color;
      if (this->left) {
        this->left->parent = this;
        this->left->restore_ordering();
      }
      if (this->right) {
        this->right->parent = this;
        this->right->restore_ordering();
      }
      this->map = rhs.map;
      return *this;
    }
    Node *child(Direction dir) {
      switch (dir) {
      case Direction::Left:
        return this->left.get();
      case Direction::Right:
        return this->right.get();
      default:
        assert(false);
      }
    }
    std::unique_ptr<Node> uchild(Direction dir) {
      switch (dir) {
      case Direction::Left:
        return std::move(this->left);
      case Direction::Right:
        return std::move(this->right);
      default:
        assert(false);
      }
    }
    std::unique_ptr<Node> uchild(Node *child) {
      return this->uchild(this->which_child(child));
    }
    std::unique_ptr<Node> &set_child(Direction dir,
                                     std::unique_ptr<Node> new_child) {

      if (new_child) {
        new_child->parent = this;
      }

      switch (dir) {
      case Direction::Left:
        return this->left = std::move(new_child);
      case Direction::Right:
        return this->right = std::move(new_child);
      default:
        assert(false);
      }
    }
    Direction which_child(Node *n) {
      if (this->left.get() == n) {
        return Direction::Left;
      }
      if (this->right.get() == n) {
        return Direction::Right;
      }
      assert(false);
    }
    void erase_child(Node *n) { this->erase_child(this->which_child(n)); }
    void erase_child(Direction dir) {
      // bringing ownership to this function scope so Deleter gets called at end
      // of function and we can do reordering things
      std::unique_ptr<Node> dropping;
      switch (dir) {
      case Direction::Right:
        dropping = std::move(this->right);
        break;
      case Direction::Left:
        dropping = std::move(this->left);
        break;
      default:
        assert(false);
      }

      // intuitively should be correct but might need to do restore ordering on
      // both instead
      dropping->prev->next = dropping->next;
      dropping->next->prev = dropping->prev;
    }
    void restore_ordering() {
      this->prev = this->calc_pred();
      this->next = this->calc_succ();
      if (this->prev) {
        this->prev->next = this;
      }
      if (this->next) {
        this->next->prev = this;
      }
    }
    Node *calc_pred() {
      if (this->left) {
        Node *ret = this->left.get();
        while (ret->right) {
          ret = ret->right.get();
        }
        return ret;
      } else {
        Node *ret = this->parent;
        Node *prev_ret = this;
        while (ret && ret->which_child(prev_ret) != Direction::Right) {
          prev_ret = ret;
          ret = prev_ret->parent;
        }
        return ret;
      }
    }
    Node *calc_succ() {
      if (this->right) {
        Node *ret = this->right.get();
        while (ret->left) {
          ret = ret->left.get();
        }
        return ret;
      } else {
        Node *ret = this->parent;
        Node *prev_ret = this;
        while (ret && ret->which_child(prev_ret) != Direction::Left) {
          prev_ret = ret;
          ret = prev_ret->parent;
        }
        return ret;
      }
    }
    void rotate(Direction dir) {

      // cannot rotate nullptr
      assert(this != nullptr);

      // we can't be root for this rotate operation
      assert(this->parent != nullptr);

      // if we're missing the child on the opposite direction this is an invalid
      // rotation
      assert(this->child(!dir));

      // gotta pull outselves out of parent to avoid accidentally overwriting
      // outselves
      std::unique_ptr<Node> self = this->parent->uchild(!dir);

      // make sure this is actually us
      assert(self.get() == this);

      // make our former position the position of the relevant child
      this->parent->set_child(!dir, this->uchild(!dir));

      // steal our former child's child
      this->set_child(!dir, this->parent->child(!dir)->uchild(dir));

      // make ourselves our former child's child
      this->parent->child(!dir)->set_child(dir, std::move(self));
    }
    // TODO:
    // Referencing
    // https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Notes_to_the_insert_diagrams
    void restore_red_black_insert(Direction dir) {
      Node *self = this;

      // infinite loop for case 2's sake, if tail recursion optimization was
      // guaranteed I'd use tail recursion
      while (true) {
        Node *parent = self->parent;
        // we're root, no-op (case 3)
        if (!parent) {
          return;
        }

        // if this is violated it's a bug
        assert(parent->child(dir) == self);

        // parent is black so no violation no-op (case 1)
        if (parent->color == Color::Black) {
          return;
        }

        Node *grandparent = parent->parent;

        // parent is root (case 4)
        if (!grandparent) {
          parent->color = Color::Black;
          return;
        }

        // table showing transforms on wikipedia doesn't have this so if it
        // happens it's probably a bug
        assert(grandparent->color == Color::Black);

        Node *uncle = grandparent->child(!grandparent->which_child(parent));
        if (uncle == nullptr || uncle->color == Color::Black) {
          if (parent->which_child(self) != grandparent->which_child(parent)) {
            // we're an inner child
            // case 5
            parent->rotate(dir);
            self = parent;
            parent = self->parent;
          }
          // case 6

          // recolor first so we aren't recoloring a dropped reference or smth
          parent->color = Color::Black;
          grandparent->color = Color::Red;
          if (grandparent->parent == nullptr) {
            map->rotate_root(!dir);
          } else {
            grandparent->rotate(!dir);
          }
          return;
        }

        // case 2 (by process of elimination)
        parent->color = Color::Black;
        uncle->color = Color::Black;
        grandparent->color = Color::Red;
        self = grandparent;
      }
    }
    // TODO:
    void restore_red_black_erase() {}
  };

  // data needed for implementation
  std::optional<Node> root;
  std::size_t _size;
  Node *min;
  Node *max;

public:
  friend Node;
  // public type definitions
  class Iterator {
    Node *underlying;
    Node *store;
    Iterator(Node *ptr, Node *potential = nullptr)
        : underlying{ptr}, store{potential} {}

  public:
    friend Map;
    Iterator() = delete;
  };
  Map() : root{}, _size{0} {}
  Map(const Map &rhs) : root{rhs.root}, _size{rhs._size} {}
  Map(Map &&rhs) : root{std::move(rhs.root)}, _size{rhs._size} {}
  Map &operator=(const Map &rhs) {
    this->root = rhs.root;
    this->_size = rhs._size;
    return *this;
  }
  Map &operator=(Map &&rhs) {
    this->root = std::move(rhs.root);
    this->_size = rhs._size;
    return *this;
  }
  std::size_t size() { return this->_size; }

private:
  // private helpers
  void rotate_root(Direction dir) {
    assert(root.has_value());

    std::unique_ptr<Node> new_root = root.value().uchild(!dir);
    // can't make null the new root
    assert(new_root);

    std::unique_ptr<Node> old_root =
        std::make_unique<Node>(std::move(root.value()));

    root.value() = std::move(*new_root);

    old_root->set_child(!dir, std::move(root.value().uchild(dir)));
    root.value().set_child(dir, std::move(old_root));
  }
  template <bool trace = false>
  std::pair<Node *, Direction> locate(const Key_T &key) {
    Node *ret_parent;
    Direction ret_dir;
    // map is empty
    if (!this->root.has_value()) {
      if constexpr (trace) {
        std::cerr << "(map empty)" << std::endl;
      }
      return std::make_pair(nullptr, ret_dir);
    }
    if constexpr (trace) {
      std::cerr << "root";
    }
    // value is in root
    if (this->root.value().val.first == key) {

      if constexpr (trace) {
        std::cerr << "->found" << std::endl;
      }

      return std::make_pair(nullptr, ret_dir);
    }
    ret_parent = &this->root.value();
    if (key < ret_parent->val.first) {
      if constexpr (trace) {
        std::cerr << "->left";
      }
      ret_dir = Direction::Left;
    } else {
      if constexpr (trace) {
        std::cerr << "->right";
      }
      ret_dir = Direction::Right;
    }
    while (ret_parent->child(ret_dir) &&
           ret_parent->child(ret_dir)->val.first != key) {
      ret_parent = ret_parent->child(ret_dir);
      if (key < ret_parent->val.first) {
        if constexpr (trace) {
          std::cerr << "->left";
        }
        ret_dir = Direction::Left;
      } else {
        if constexpr (trace) {
          std::cerr << "->right";
        }
        ret_dir = Direction::Right;
      }
    }
    if constexpr (trace) {
      std::cerr << "->found" << std::endl;
    }
    return std::make_pair(ret_parent, ret_dir);
  }

public:
  // baseline find using locate
  Iterator find(const Key_T &key) {
    auto [parent, dir] = locate(key);
    if (parent == nullptr) {
      if (this->root.has_value()) {
        if (this->root.value().val.first == key) {
          return Iterator{&this->root.value()};
        }
      }
      return this->end();
    }
    if (parent->child(dir) != nullptr) {
      return Iterator{parent->child(dir), nullptr};
    }
    return this->end();
  }
  Iterator find_trace(const Key_T &key) {
    auto [parent, dir] = locate<true>(key);
    if (parent == nullptr) {
      if (this->root.has_value()) {
        if (this->root.value().val.first == key) {
          return Iterator{&this->root.value()};
        }
      }
      return this->end();
    }
    if (parent->child(dir) != nullptr) {
      return Iterator{parent->child(dir), nullptr};
    }
    return this->end();
  }
  // baseline modification operations
  std::pair<Iterator, bool> insert(const ValueType &val) {
    // for convenience
    auto &[key, map] = val;

    auto [parent, dir] = locate(key);
    // located root node
    if (parent == nullptr) {
      if (this->root.has_value()) {
        return std::make_pair(Iterator{&root.value()}, false);
      } else {
        this->root = Node{val, this};
        return std::make_pair(Iterator{&root.value()}, true);
      }
    }

    // non-root node
    if (parent->child(dir)) {
      // node already present
      return std::make_pair(Iterator{parent->child(dir)}, false);
    }

    // need to insert non-root node
    Node *new_node =
        parent->set_child(dir, std::make_unique<Node>(Node{val, this})).get();
    new_node->restore_red_black_insert(dir);
    new_node->restore_ordering();
    return std::make_pair(Iterator{new_node}, true);
  }
  void erase(Iterator pos) {
    Node *p = pos.underlying->parent;
    p->erase_child(pos.underlying);
    p->restore_red_black_erase();
  }

  // baseline iterator creation
  Iterator begin() { return Iterator{min, nullptr}; }
  Iterator end() { return Iterator{nullptr, max}; }
};
} // namespace cs440
#endif