michael@0: /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 4 -*-
michael@0:  * vim: set ts=8 sts=4 et sw=4 tw=99:
michael@0:  * This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this
michael@0:  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: #ifndef ds_SplayTree_h
michael@0: #define ds_SplayTree_h
michael@0: 
michael@0: #include "ds/LifoAlloc.h"
michael@0: 
michael@0: namespace js {
michael@0: 
michael@0: /*
michael@0:  * Class which represents a splay tree with nodes allocated from a LifoAlloc.
michael@0:  * Splay trees are balanced binary search trees for which search, insert and
michael@0:  * remove are all amortized O(log n).
michael@0:  *
michael@0:  * T indicates the type of tree elements, C must have a static
michael@0:  * compare(const T&, const T&) method ordering the elements. As for LifoAlloc
michael@0:  * objects, T objects stored in the tree will not be explicitly destroyed.
michael@0:  */
michael@0: template <class T, class C>
michael@0: class SplayTree
michael@0: {
michael@0:     struct Node {
michael@0:         T item;
michael@0:         Node *left, *right, *parent;
michael@0: 
michael@0:         Node(const T &item)
michael@0:           : item(item), left(nullptr), right(nullptr), parent(nullptr)
michael@0:         {}
michael@0:     };
michael@0: 
michael@0:     LifoAlloc *alloc;
michael@0:     Node *root, *freeList;
michael@0: 
michael@0:     SplayTree(const SplayTree &) MOZ_DELETE;
michael@0:     SplayTree &operator=(const SplayTree &) MOZ_DELETE;
michael@0: 
michael@0:   public:
michael@0: 
michael@0:     SplayTree(LifoAlloc *alloc = nullptr)
michael@0:       : alloc(alloc), root(nullptr), freeList(nullptr)
michael@0:     {}
michael@0: 
michael@0:     void setAllocator(LifoAlloc *alloc) {
michael@0:         this->alloc = alloc;
michael@0:     }
michael@0: 
michael@0:     bool empty() const {
michael@0:         return !root;
michael@0:     }
michael@0: 
michael@0:     bool contains(const T &v, T *res)
michael@0:     {
michael@0:         if (!root)
michael@0:             return false;
michael@0:         Node *last = lookup(v);
michael@0:         splay(last);
michael@0:         checkCoherency(root, nullptr);
michael@0:         if (C::compare(v, last->item) == 0) {
michael@0:             *res = last->item;
michael@0:             return true;
michael@0:         }
michael@0:         return false;
michael@0:     }
michael@0: 
michael@0:     bool insert(const T &v)
michael@0:     {
michael@0:         Node *element = allocateNode(v);
michael@0:         if (!element)
michael@0:             return false;
michael@0: 
michael@0:         if (!root) {
michael@0:             root = element;
michael@0:             return true;
michael@0:         }
michael@0:         Node *last = lookup(v);
michael@0:         int cmp = C::compare(v, last->item);
michael@0: 
michael@0:         // Don't tolerate duplicate elements.
michael@0:         JS_ASSERT(cmp);
michael@0: 
michael@0:         Node *&parentPointer = (cmp < 0) ? last->left : last->right;
michael@0:         JS_ASSERT(!parentPointer);
michael@0:         parentPointer = element;
michael@0:         element->parent = last;
michael@0: 
michael@0:         splay(element);
michael@0:         checkCoherency(root, nullptr);
michael@0:         return true;
michael@0:     }
michael@0: 
michael@0:     void remove(const T &v)
michael@0:     {
michael@0:         Node *last = lookup(v);
michael@0:         JS_ASSERT(last && C::compare(v, last->item) == 0);
michael@0: 
michael@0:         splay(last);
michael@0:         JS_ASSERT(last == root);
michael@0: 
michael@0:         // Find another node which can be swapped in for the root: either the
michael@0:         // rightmost child of the root's left, or the leftmost child of the
michael@0:         // root's right.
michael@0:         Node *swap, *swapChild;
michael@0:         if (root->left) {
michael@0:             swap = root->left;
michael@0:             while (swap->right)
michael@0:                 swap = swap->right;
michael@0:             swapChild = swap->left;
michael@0:         } else if (root->right) {
michael@0:             swap = root->right;
michael@0:             while (swap->left)
michael@0:                 swap = swap->left;
michael@0:             swapChild = swap->right;
michael@0:         } else {
michael@0:             freeNode(root);
michael@0:             root = nullptr;
michael@0:             return;
michael@0:         }
michael@0: 
michael@0:         // The selected node has at most one child, in swapChild. Detach it
michael@0:         // from the subtree by replacing it with that child.
michael@0:         if (swap == swap->parent->left)
michael@0:             swap->parent->left = swapChild;
michael@0:         else
michael@0:             swap->parent->right = swapChild;
michael@0:         if (swapChild)
michael@0:             swapChild->parent = swap->parent;
michael@0: 
michael@0:         root->item = swap->item;
michael@0:         freeNode(swap);
michael@0: 
michael@0:         checkCoherency(root, nullptr);
michael@0:     }
michael@0: 
michael@0:     template <class Op>
michael@0:     void forEach(Op op)
michael@0:     {
michael@0:         forEachInner(op, root);
michael@0:     }
michael@0: 
michael@0:   private:
michael@0: 
michael@0:     Node *lookup(const T &v)
michael@0:     {
michael@0:         JS_ASSERT(root);
michael@0:         Node *node = root, *parent;
michael@0:         do {
michael@0:             parent = node;
michael@0:             int c = C::compare(v, node->item);
michael@0:             if (c == 0)
michael@0:                 return node;
michael@0:             else if (c < 0)
michael@0:                 node = node->left;
michael@0:             else
michael@0:                 node = node->right;
michael@0:         } while (node);
michael@0:         return parent;
michael@0:     }
michael@0: 
michael@0:     Node *allocateNode(const T &v)
michael@0:     {
michael@0:         Node *node = freeList;
michael@0:         if (node) {
michael@0:             freeList = node->left;
michael@0:             new(node) Node(v);
michael@0:             return node;
michael@0:         }
michael@0:         return alloc->new_<Node>(v);
michael@0:     }
michael@0: 
michael@0:     void freeNode(Node *node)
michael@0:     {
michael@0:         node->left = freeList;
michael@0:         freeList = node;
michael@0:     }
michael@0: 
michael@0:     void splay(Node *node)
michael@0:     {
michael@0:         // Rotate the element until it is at the root of the tree. Performing
michael@0:         // the rotations in this fashion preserves the amortized balancing of
michael@0:         // the tree.
michael@0:         JS_ASSERT(node);
michael@0:         while (node != root) {
michael@0:             Node *parent = node->parent;
michael@0:             if (parent == root) {
michael@0:                 // Zig rotation.
michael@0:                 rotate(node);
michael@0:                 JS_ASSERT(node == root);
michael@0:                 return;
michael@0:             }
michael@0:             Node *grandparent = parent->parent;
michael@0:             if ((parent->left == node) == (grandparent->left == parent)) {
michael@0:                 // Zig-zig rotation.
michael@0:                 rotate(parent);
michael@0:                 rotate(node);
michael@0:             } else {
michael@0:                 // Zig-zag rotation.
michael@0:                 rotate(node);
michael@0:                 rotate(node);
michael@0:             }
michael@0:         }
michael@0:     }
michael@0: 
michael@0:     void rotate(Node *node)
michael@0:     {
michael@0:         // Rearrange nodes so that node becomes the parent of its current
michael@0:         // parent, while preserving the sortedness of the tree.
michael@0:         Node *parent = node->parent;
michael@0:         if (parent->left == node) {
michael@0:             //     x          y
michael@0:             //   y  c  ==>  a  x
michael@0:             //  a b           b c
michael@0:             parent->left = node->right;
michael@0:             if (node->right)
michael@0:                 node->right->parent = parent;
michael@0:             node->right = parent;
michael@0:         } else {
michael@0:             JS_ASSERT(parent->right == node);
michael@0:             //   x             y
michael@0:             //  a  y   ==>   x  c
michael@0:             //    b c       a b
michael@0:             parent->right = node->left;
michael@0:             if (node->left)
michael@0:                 node->left->parent = parent;
michael@0:             node->left = parent;
michael@0:         }
michael@0:         node->parent = parent->parent;
michael@0:         parent->parent = node;
michael@0:         if (Node *grandparent = node->parent) {
michael@0:             if (grandparent->left == parent)
michael@0:                 grandparent->left = node;
michael@0:             else
michael@0:                 grandparent->right = node;
michael@0:         } else {
michael@0:             root = node;
michael@0:         }
michael@0:     }
michael@0: 
michael@0:     template <class Op>
michael@0:     void forEachInner(Op op, Node *node)
michael@0:     {
michael@0:         if (!node)
michael@0:             return;
michael@0: 
michael@0:         forEachInner(op, node->left);
michael@0:         op(node->item);
michael@0:         forEachInner(op, node->right);
michael@0:     }
michael@0: 
michael@0:     Node *checkCoherency(Node *node, Node *minimum)
michael@0:     {
michael@0: #ifdef DEBUG
michael@0:         if (!node) {
michael@0:             JS_ASSERT(!root);
michael@0:             return nullptr;
michael@0:         }
michael@0:         JS_ASSERT_IF(!node->parent, node == root);
michael@0:         JS_ASSERT_IF(minimum, C::compare(minimum->item, node->item) < 0);
michael@0:         if (node->left) {
michael@0:             JS_ASSERT(node->left->parent == node);
michael@0:             Node *leftMaximum = checkCoherency(node->left, minimum);
michael@0:             JS_ASSERT(C::compare(leftMaximum->item, node->item) < 0);
michael@0:         }
michael@0:         if (node->right) {
michael@0:             JS_ASSERT(node->right->parent == node);
michael@0:             return checkCoherency(node->right, node);
michael@0:         }
michael@0:         return node;
michael@0: #else
michael@0:         return nullptr;
michael@0: #endif
michael@0:     }
michael@0: };
michael@0: 
michael@0: }  /* namespace js */
michael@0: 
michael@0: #endif /* ds_SplayTree_h */