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