mfbt/SplayTree.h

   1 /* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
   2 /* vim: set ts=8 sts=2 et sw=2 tw=80: */
   3 /* This Source Code Form is subject to the terms of the Mozilla Public
   4  * License, v. 2.0. If a copy of the MPL was not distributed with this
   5  * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
   6
   7 /**
   8  * A sorted tree with optimal access times, where recently-accessed elements
   9  * are faster to access again.
  10  */
  11
  12 #ifndef mozilla_SplayTree_h
  13 #define mozilla_SplayTree_h
  14
  15 #include "mozilla/Assertions.h"
  16 #include "mozilla/Attributes.h"
  17
  18 namespace mozilla {
  19
  20 template<class T, class C>
  21 class SplayTree;
  22
  23 template<typename T>
  24 class SplayTreeNode
  25 {
  26 public:
  27   template<class A, class B>
  28   friend class SplayTree;
  29
  30   SplayTreeNode()
  31     : mLeft(nullptr)
  32     , mRight(nullptr)
  33     , mParent(nullptr)
  34   {}
  35
  36 private:
  37   T* mLeft;
  38   T* mRight;
  39   T* mParent;
  40 };
  41
  42
  43 /**
  44  * Class which represents a splay tree.
  45  * Splay trees are balanced binary search trees for which search, insert and
  46  * remove are all amortized O(log n), but where accessing a node makes it
  47  * faster to access that node in the future.
  48  *
  49  * T indicates the type of tree elements, Comparator must have a static
  50  * compare(const T&, const T&) method ordering the elements. The compare
  51  * method must be free from side effects.
  52  */
  53 template<typename T, class Comparator>
  54 class SplayTree
  55 {
  56   T* mRoot;
  57
  58 public:
  59   constexpr SplayTree()
  60     : mRoot(nullptr)
  61   {}
  62
  63   bool empty() const
  64   {
  65     return !mRoot;
  66   }
  67
  68   T* find(const T& aValue)
  69   {
  70     if (empty()) {
  71       return nullptr;
  72     }
  73
  74     T* last = lookup(aValue);
  75     splay(last);
  76     return Comparator::compare(aValue, *last) == 0 ? last : nullptr;
  77   }
  78
  79   void insert(T* aValue)
  80   {
  81     MOZ_ASSERT(!find(*aValue), "Duplicate elements are not allowed.");
  82
  83     if (!mRoot) {
  84       mRoot = aValue;
  85       return;
  86     }
  87     T* last = lookup(*aValue);
  88     int cmp = Comparator::compare(*aValue, *last);
  89
  90     finishInsertion(last, cmp, aValue);
  91   }
  92
  93   T* findOrInsert(const T& aValue);
  94
  95   T* remove(const T& aValue)
  96   {
  97     T* last = lookup(aValue);
  98     MOZ_ASSERT(last, "This tree must contain the element being removed.");
  99     MOZ_ASSERT(Comparator::compare(aValue, *last) == 0);
 100
 101     // Splay the tree so that the item to remove is the root.
 102     splay(last);
 103     MOZ_ASSERT(last == mRoot);
 104
 105     // Find another node which can be swapped in for the root: either the
 106     // rightmost child of the root's left, or the leftmost child of the
 107     // root's right.
 108     T* swap;
 109     T* swapChild;
 110     if (mRoot->mLeft) {
 111       swap = mRoot->mLeft;
 112       while (swap->mRight) {
 113         swap = swap->mRight;
 114       }
 115       swapChild = swap->mLeft;
 116     } else if (mRoot->mRight) {
 117       swap = mRoot->mRight;
 118       while (swap->mLeft) {
 119         swap = swap->mLeft;
 120       }
 121       swapChild = swap->mRight;
 122     } else {
 123       T* result = mRoot;
 124       mRoot = nullptr;
 125       return result;
 126     }
 127
 128     // The selected node has at most one child, in swapChild. Detach it
 129     // from the subtree by replacing it with that child.
 130     if (swap == swap->mParent->mLeft) {
 131       swap->mParent->mLeft = swapChild;
 132     } else {
 133       swap->mParent->mRight = swapChild;
 134     }
 135     if (swapChild) {
 136       swapChild->mParent = swap->mParent;
 137     }
 138
 139     // Make the selected node the new root.
 140     mRoot = swap;
 141     mRoot->mParent = nullptr;
 142     mRoot->mLeft = last->mLeft;
 143     mRoot->mRight = last->mRight;
 144     if (mRoot->mLeft) {
 145       mRoot->mLeft->mParent = mRoot;
 146     }
 147     if (mRoot->mRight) {
 148       mRoot->mRight->mParent = mRoot;
 149     }
 150
 151     last->mLeft = nullptr;
 152     last->mRight = nullptr;
 153     return last;
 154   }
 155
 156   T* removeMin()
 157   {
 158     MOZ_ASSERT(mRoot, "No min to remove!");
 159
 160     T* min = mRoot;
 161     while (min->mLeft) {
 162       min = min->mLeft;
 163     }
 164     return remove(*min);
 165   }
 166
 167   // For testing purposes only.
 168   void checkCoherency()
 169   {
 170     checkCoherency(mRoot, nullptr);
 171   }
 172
 173 private:
 174   /**
 175    * Returns the node in this comparing equal to |aValue|, or a node just
 176    * greater or just less than |aValue| if there is no such node.
 177    */
 178   T* lookup(const T& aValue)
 179   {
 180     MOZ_ASSERT(!empty());
 181
 182     T* node = mRoot;
 183     T* parent;
 184     do {
 185       parent = node;
 186       int c = Comparator::compare(aValue, *node);
 187       if (c == 0) {
 188         return node;
 189       } else if (c < 0) {
 190         node = node->mLeft;
 191       } else {
 192         node = node->mRight;
 193       }
 194     } while (node);
 195     return parent;
 196   }
 197
 198   void finishInsertion(T* aLast, int32_t aCmp, T* aNew)
 199   {
 200     MOZ_ASSERT(aCmp, "Nodes shouldn't be equal!");
 201
 202     T** parentPointer = (aCmp < 0) ? &aLast->mLeft : &aLast->mRight;
 203     MOZ_ASSERT(!*parentPointer);
 204     *parentPointer = aNew;
 205     aNew->mParent = aLast;
 206
 207     splay(aNew);
 208   }
 209
 210   /**
 211    * Rotate the tree until |node| is at the root of the tree. Performing
 212    * the rotations in this fashion preserves the amortized balancing of
 213    * the tree.
 214    */
 215   void splay(T* aNode)
 216   {
 217     MOZ_ASSERT(aNode);
 218
 219     while (aNode != mRoot) {
 220       T* parent = aNode->mParent;
 221       if (parent == mRoot) {
 222         // Zig rotation.
 223         rotate(aNode);
 224         MOZ_ASSERT(aNode == mRoot);
 225         return;
 226       }
 227       T* grandparent = parent->mParent;
 228       if ((parent->mLeft == aNode) == (grandparent->mLeft == parent)) {
 229         // Zig-zig rotation.
 230         rotate(parent);
 231         rotate(aNode);
 232       } else {
 233         // Zig-zag rotation.
 234         rotate(aNode);
 235         rotate(aNode);
 236       }
 237     }
 238   }
 239
 240   void rotate(T* aNode)
 241   {
 242     // Rearrange nodes so that aNode becomes the parent of its current
 243     // parent, while preserving the sortedness of the tree.
 244     T* parent = aNode->mParent;
 245     if (parent->mLeft == aNode) {
 246       //     x          y
 247       //   y  c  ==>  a  x
 248       //  a b           b c
 249       parent->mLeft = aNode->mRight;
 250       if (aNode->mRight) {
 251         aNode->mRight->mParent = parent;
 252       }
 253       aNode->mRight = parent;
 254     } else {
 255       MOZ_ASSERT(parent->mRight == aNode);
 256       //   x             y
 257       //  a  y   ==>   x  c
 258       //    b c       a b
 259       parent->mRight = aNode->mLeft;
 260       if (aNode->mLeft) {
 261         aNode->mLeft->mParent = parent;
 262       }
 263       aNode->mLeft = parent;
 264     }
 265     aNode->mParent = parent->mParent;
 266     parent->mParent = aNode;
 267     if (T* grandparent = aNode->mParent) {
 268       if (grandparent->mLeft == parent) {
 269         grandparent->mLeft = aNode;
 270       } else {
 271         grandparent->mRight = aNode;
 272       }
 273     } else {
 274       mRoot = aNode;
 275     }
 276   }
 277
 278   T* checkCoherency(T* aNode, T* aMinimum)
 279   {
 280     if (mRoot) {
 281       MOZ_RELEASE_ASSERT(!mRoot->mParent);
 282     }
 283     if (!aNode) {
 284       MOZ_RELEASE_ASSERT(!mRoot);
 285       return nullptr;
 286     }
 287     if (!aNode->mParent) {
 288       MOZ_RELEASE_ASSERT(aNode == mRoot);
 289     }
 290     if (aMinimum) {
 291       MOZ_RELEASE_ASSERT(Comparator::compare(*aMinimum, *aNode) < 0);
 292     }
 293     if (aNode->mLeft) {
 294       MOZ_RELEASE_ASSERT(aNode->mLeft->mParent == aNode);
 295       T* leftMaximum = checkCoherency(aNode->mLeft, aMinimum);
 296       MOZ_RELEASE_ASSERT(Comparator::compare(*leftMaximum, *aNode) < 0);
 297     }
 298     if (aNode->mRight) {
 299       MOZ_RELEASE_ASSERT(aNode->mRight->mParent == aNode);
 300       return checkCoherency(aNode->mRight, aNode);
 301     }
 302     return aNode;
 303   }
 304
 305   SplayTree(const SplayTree&) = delete;
 306   void operator=(const SplayTree&) = delete;
 307 };
 308
 309 template<typename T, class Comparator>
 310 T*
 311 SplayTree<T, Comparator>::findOrInsert(const T& aValue)
 312 {
 313   if (!mRoot) {
 314     mRoot = new T(aValue);
 315     return mRoot;
 316   }
 317
 318   T* last = lookup(aValue);
 319   int cmp = Comparator::compare(aValue, *last);
 320   if (!cmp) {
 321     return last;
 322   }
 323
 324   T* t = new T(aValue);
 325   finishInsertion(last, cmp, t);
 326   return t;
 327 }
 328
 329 }  /* namespace mozilla */
 330
 331 #endif /* mozilla_SplayTree_h */