/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */

/* vim: set ts=8 sts=2 et sw=2 tw=80: */

/* This Source Code Form is subject to the terms of the Mozilla Public

 * License, v. 2.0. If a copy of the MPL was not distributed with this

 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/**

 * A sorted tree with optimal access times, where recently-accessed elements

 * are faster to access again.

 */

#ifndef mozilla_SplayTree_h

#define mozilla_SplayTree_h

#include "mozilla/Assertions.h"

#include "mozilla/Attributes.h"

namespace mozilla {

template<class T, class C>

class SplayTree;

template<typename T>

class SplayTreeNode

{

public:

  template<class A, class B>

  friend class SplayTree;

  SplayTreeNode()

    : mLeft(nullptr)

    , mRight(nullptr)

    , mParent(nullptr)

  {}

private:

  T* mLeft;

  T* mRight;

  T* mParent;

};

/**

 * Class which represents a splay tree.

 * Splay trees are balanced binary search trees for which search, insert and

 * remove are all amortized O(log n), but where accessing a node makes it

 * faster to access that node in the future.

 *

 * T indicates the type of tree elements, Comparator must have a static

 * compare(const T&, const T&) method ordering the elements. The compare

 * method must be free from side effects.

 */

template<typename T, class Comparator>

class SplayTree

{

  T* mRoot;

public:

  constexpr SplayTree()

    : mRoot(nullptr)

  {}

  bool empty() const

  {

    return !mRoot;

  }

  T* find(const T& aValue)

  {

    if (empty()) {

      return nullptr;

    }

    T* last = lookup(aValue);

    splay(last);

    return Comparator::compare(aValue, *last) == 0 ? last : nullptr;

  }

  void insert(T* aValue)

  {

    MOZ_ASSERT(!find(*aValue), "Duplicate elements are not allowed.");

    if (!mRoot) {

      mRoot = aValue;

      return;

    }

    T* last = lookup(*aValue);

    int cmp = Comparator::compare(*aValue, *last);

    finishInsertion(last, cmp, aValue);

  }

  T* findOrInsert(const T& aValue);

  T* remove(const T& aValue)

  {

    T* last = lookup(aValue);

    MOZ_ASSERT(last, "This tree must contain the element being removed.");

    MOZ_ASSERT(Comparator::compare(aValue, *last) == 0);

    // Splay the tree so that the item to remove is the root.

    splay(last);

    MOZ_ASSERT(last == mRoot);

    // Find another node which can be swapped in for the root: either the

    // rightmost child of the root's left, or the leftmost child of the

    // root's right.

    T* swap;

    T* swapChild;

    if (mRoot->mLeft) {

      swap = mRoot->mLeft;

      while (swap->mRight) {

        swap = swap->mRight;

      }

      swapChild = swap->mLeft;

    } else if (mRoot->mRight) {

      swap = mRoot->mRight;

      while (swap->mLeft) {

        swap = swap->mLeft;

      }

      swapChild = swap->mRight;

    } else {

      T* result = mRoot;

      mRoot = nullptr;

      return result;

    }

    // The selected node has at most one child, in swapChild. Detach it

    // from the subtree by replacing it with that child.

    if (swap == swap->mParent->mLeft) {

      swap->mParent->mLeft = swapChild;

    } else {

      swap->mParent->mRight = swapChild;

    }

    if (swapChild) {

      swapChild->mParent = swap->mParent;

    }

    // Make the selected node the new root.

    mRoot = swap;

    mRoot->mParent = nullptr;

    mRoot->mLeft = last->mLeft;

    mRoot->mRight = last->mRight;

    if (mRoot->mLeft) {

      mRoot->mLeft->mParent = mRoot;

    }

    if (mRoot->mRight) {

      mRoot->mRight->mParent = mRoot;

    }

    last->mLeft = nullptr;

    last->mRight = nullptr;

    return last;

  }

  T* removeMin()

  {

    MOZ_ASSERT(mRoot, "No min to remove!");

    T* min = mRoot;

    while (min->mLeft) {

      min = min->mLeft;

    }

    return remove(*min);

  }

  // For testing purposes only.

  void checkCoherency()

  {

    checkCoherency(mRoot, nullptr);

  }

private:

  /**

   * Returns the node in this comparing equal to |aValue|, or a node just

   * greater or just less than |aValue| if there is no such node.

   */

  T* lookup(const T& aValue)

  {

    MOZ_ASSERT(!empty());

    T* node = mRoot;

    T* parent;

    do {

      parent = node;

      int c = Comparator::compare(aValue, *node);

      if (c == 0) {

        return node;

      } else if (c < 0) {

        node = node->mLeft;

      } else {

        node = node->mRight;

      }

    } while (node);

    return parent;

  }

  void finishInsertion(T* aLast, int32_t aCmp, T* aNew)

  {

    MOZ_ASSERT(aCmp, "Nodes shouldn't be equal!");

    T** parentPointer = (aCmp < 0) ? &aLast->mLeft : &aLast->mRight;

    MOZ_ASSERT(!*parentPointer);

    *parentPointer = aNew;

    aNew->mParent = aLast;

    splay(aNew);

  }

  /**

   * Rotate the tree until |node| is at the root of the tree. Performing

   * the rotations in this fashion preserves the amortized balancing of

   * the tree.

   */

  void splay(T* aNode)

  {

    MOZ_ASSERT(aNode);

    while (aNode != mRoot) {

      T* parent = aNode->mParent;

      if (parent == mRoot) {

        // Zig rotation.

        rotate(aNode);

        MOZ_ASSERT(aNode == mRoot);

        return;

      }

      T* grandparent = parent->mParent;

      if ((parent->mLeft == aNode) == (grandparent->mLeft == parent)) {

        // Zig-zig rotation.

        rotate(parent);

        rotate(aNode);

      } else {

        // Zig-zag rotation.

        rotate(aNode);

        rotate(aNode);

      }

    }

  }

  void rotate(T* aNode)

  {

    // Rearrange nodes so that aNode becomes the parent of its current

    // parent, while preserving the sortedness of the tree.

    T* parent = aNode->mParent;

    if (parent->mLeft == aNode) {

      //     x          y

      //   y  c  ==>  a  x

      //  a b           b c

      parent->mLeft = aNode->mRight;

      if (aNode->mRight) {

        aNode->mRight->mParent = parent;

      }

      aNode->mRight = parent;

    } else {

      MOZ_ASSERT(parent->mRight == aNode);

      //   x             y

      //  a  y   ==>   x  c

      //    b c       a b

      parent->mRight = aNode->mLeft;

      if (aNode->mLeft) {

        aNode->mLeft->mParent = parent;

      }

      aNode->mLeft = parent;

    }

    aNode->mParent = parent->mParent;

    parent->mParent = aNode;

    if (T* grandparent = aNode->mParent) {

      if (grandparent->mLeft == parent) {

        grandparent->mLeft = aNode;

      } else {

        grandparent->mRight = aNode;

      }

    } else {

      mRoot = aNode;

    }

  }

  T* checkCoherency(T* aNode, T* aMinimum)

  {

    if (mRoot) {

      MOZ_RELEASE_ASSERT(!mRoot->mParent);

    }

    if (!aNode) {

      MOZ_RELEASE_ASSERT(!mRoot);

      return nullptr;

    }

    if (!aNode->mParent) {

      MOZ_RELEASE_ASSERT(aNode == mRoot);

    }

    if (aMinimum) {

      MOZ_RELEASE_ASSERT(Comparator::compare(*aMinimum, *aNode) < 0);

    }

    if (aNode->mLeft) {

      MOZ_RELEASE_ASSERT(aNode->mLeft->mParent == aNode);

      T* leftMaximum = checkCoherency(aNode->mLeft, aMinimum);

      MOZ_RELEASE_ASSERT(Comparator::compare(*leftMaximum, *aNode) < 0);

    }

    if (aNode->mRight) {

      MOZ_RELEASE_ASSERT(aNode->mRight->mParent == aNode);

      return checkCoherency(aNode->mRight, aNode);

    }

    return aNode;

  }

  SplayTree(const SplayTree&) = delete;

  void operator=(const SplayTree&) = delete;

};

template<typename T, class Comparator>

T*

SplayTree<T, Comparator>::findOrInsert(const T& aValue)

{

  if (!mRoot) {

    mRoot = new T(aValue);

    return mRoot;

  }

  T* last = lookup(aValue);

  int cmp = Comparator::compare(aValue, *last);

  if (!cmp) {

    return last;

  }

  T* t = new T(aValue);

  finishInsertion(last, cmp, t);

  return t;

}

}  /* namespace mozilla */

#endif /* mozilla_SplayTree_h */