Bug 537339 - Add heap functions to nsTArray, r=roc

xpcom/glue/nsTArray.h

@@ -877,7 +877,7 @@ class nsTArray : public nsTArray_base {
 
     // This method sorts the elements of the array.  It uses the LessThan
     // method defined on the given Comparator object to collate elements.
-    // @param c  The Comparator to used to collate elements.
+    // @param comp The Comparator used to collate elements.
     template<class Comparator>
     void Sort(const Comparator& comp) {
       NS_QuickSort(Elements(), Length(), sizeof(elem_type),
@@ -891,6 +891,77 @@ class nsTArray : public nsTArray_base {
       Sort(nsDefaultComparator<elem_type, elem_type>());
     }
 
+    //
+    // Binary Heap
+    //
+
+    // Sorts the array into a binary heap.
+    // @param comp The Comparator used to create the heap
+    template<class Comparator>
+    void MakeHeap(const Comparator& comp) {
+      if (!Length()) {
+        return;
+      }
+      elem_type *elem = Elements();
+      index_type index = (Length() - 1) / 2;
+      do {
+        SiftDown(index, comp);
+      } while (index--);
+    }
+
+    // A variation on the MakeHeap method defined above.
+    void MakeHeap() {
+      MakeHeap(nsDefaultComparator<elem_type, elem_type>());
+    }
+
+    // Adds an element to the heap
+    // @param item The item to add
+    // @param comp The Comparator used to sift-up the item
+    template<class Item, class Comparator>
+    elem_type *PushHeap(const Item& item, const Comparator& comp) {
+      if (!nsTArray_base::InsertSlotsAt(Length(), 1, sizeof(elem_type))) {
+        return nsnull;
+      }
+      // Sift up the new node
+      elem_type *elem = Elements();
+      index_type index = Length() - 1;
+      index_type parent_index = (index - 1) / 2;
+      while (index && comp.LessThan(elem[parent_index], item)) {
+        elem[index] = elem[parent_index];
+        index = parent_index;
+        parent_index = (index - 1) / 2;
+      }
+      elem[index] = item;
+      return &elem[index];
+    }
+
+    // A variation on the PushHeap method defined above.
+    template<class Item>
+    elem_type *PushHeap(const Item& item) {
+      return PushHeap(item, nsDefaultComparator<elem_type, Item>());
+    }
+
+    // Delete the root of the heap and restore the heap
+    // @param comp The Comparator used to restore the heap
+    template<class Comparator>
+    void PopHeap(const Comparator& comp) {
+      if (!Length()) {
+        return;
+      }
+      index_type last_index = Length() - 1;
+      elem_type *elem = Elements();
+      elem[0] = elem[last_index];
+      TruncateLength(last_index);
+      if (Length()) {
+        SiftDown(0, comp);
+      }
+    }
+
+    // A variation on the PopHeap method defined above.
+    void PopHeap() {
+      PopHeap(nsDefaultComparator<elem_type, elem_type>());
+    }
+
   protected:
 
     // This method invokes elem_type's destructor on a range of elements.
@@ -915,6 +986,36 @@ class nsTArray : public nsTArray_base {
         elem_traits::Construct(iter, *values);
       }
     }
+
+    // This method sifts an item down to its proper place in a binary heap
+    // @param index The index of the node to start sifting down from
+    // @param comp  The Comparator used to sift down
+    template<class Comparator>
+    void SiftDown(index_type index, const Comparator& comp) {
+      elem_type *elem = Elements();
+      elem_type item = elem[index];
+      index_type end = Length() - 1;
+      while ((index * 2) < end) {
+        const index_type left = (index * 2) + 1;
+        const index_type right = (index * 2) + 2;
+        const index_type parent_index = index;
+        if (comp.LessThan(item, elem[left])) {
+          if (left < end &&
+              comp.LessThan(elem[left], elem[right])) {
+            index = right;
+          } else {
+            index = left;
+          }
+        } else if (left < end &&
+                   comp.LessThan(item, elem[right])) {
+          index = right;
+        } else {
+          break;
+        }
+        elem[parent_index] = elem[index];
+      }
+      elem[index] = item;
+    }
 };
 
 template<class E, PRUint32 N>

xpcom/tests/TestTArray.cpp

@@ -523,6 +523,45 @@ static PRBool test_indexof() {
 
 //----
 
+template <class Array>
+static PRBool is_heap(const Array& ary, PRUint32 len) {
+  PRUint32 index = 1;
+  while (index < len) {
+    if (ary[index] > ary[(index - 1) >> 1])
+      return PR_FALSE;
+    index++;
+  }
+  return PR_TRUE;
+} 
+
+static PRBool test_heap() {
+  const int data[] = {4,6,8,2,4,1,5,7,3};
+  nsTArray<int> ary;
+  ary.AppendElements(data, NS_ARRAY_LENGTH(data));
+  // make a heap and make sure it's a heap
+  ary.MakeHeap();
+  if (!is_heap(ary, NS_ARRAY_LENGTH(data)))
+    return PR_FALSE;
+  // pop the root and make sure it's still a heap
+  int root = ary[0];
+  ary.PopHeap();
+  if (!is_heap(ary, NS_ARRAY_LENGTH(data) - 1))
+    return PR_FALSE;
+  // push the previously poped value back on and make sure it's still a heap
+  ary.PushHeap(root);
+  if (!is_heap(ary, NS_ARRAY_LENGTH(data)))
+    return PR_FALSE;
+  // make sure the heap looks like what we expect
+  const int expected_data[] = {8,7,5,6,4,1,4,2,3};
+  PRUint32 index;
+  for (index = 0; index < NS_ARRAY_LENGTH(data); index++)
+    if (ary[index] != expected_data[index])
+      return PR_FALSE;
+  return PR_TRUE;
+}
+
+//----
+
 typedef PRBool (*TestFunc)();
 #define DECL_TEST(name) { #name, name }
 
@@ -543,6 +582,7 @@ static const struct Test {
   DECL_TEST(test_autoarray),
 #endif
   DECL_TEST(test_indexof),
+  DECL_TEST(test_heap),
   { nsnull, nsnull }
 };