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);
  }
}
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;
    Node(internal_ValueType val)
        : parent{nullptr}, val{val}, left{}, right{}, color{Color::Red},
          prev{nullptr}, next{nullptr} {}
    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} {
      this->left->parent = this;
      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} {
      this->left->parent = this;
      this->right->parent = this;
    }
    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->left->parent = this;
      this->right->parent = this;
      this->left->restore_ordering();
      this->right->restore_ordering();
      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->left->parent = this;
      this->right->parent = this;
      this->left->restore_ordering();
      this->right->restore_ordering();
      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) {

      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();
      this->prev->next = this;
      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:
    void restore_red_black_insert() {}
    // TODO:
    void restore_red_black_erase() {}
  };

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

public:
  // 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:
  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) {
      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};
        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})).get();
    new_node->restore_red_black_insert();
    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