Bug 537339 - Add heap functions to nsTArray, r=roc
authorMichael Wu <mwu@mozilla.com>
Wed, 05 May 2010 10:21:23 -0700
changeset 41924 0c65ff82f0e9d01fc89653cbe25d241628085272
parent 41923 6c5782d2377eabc3bf262f1d3a843bdbf980baf9
child 41925 5c843ae12c40556f28f32493cbfcb4b89467ac66
push idunknown
push userunknown
push dateunknown
reviewersroc
bugs537339
milestone1.9.3a5pre
Bug 537339 - Add heap functions to nsTArray, r=roc
xpcom/glue/nsTArray.h
xpcom/tests/TestTArray.cpp
--- a/xpcom/glue/nsTArray.h
+++ b/xpcom/glue/nsTArray.h
@@ -872,30 +872,101 @@ class nsTArray : public nsTArray_base {
     }
 
     //
     // Sorting
     //
 
     // 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),
                    nsQuickSortComparator<elem_type, Comparator>::Compare,
                    const_cast<Comparator*>(&comp));
     }
 
     // A variation on the Sort method defined above that assumes that
     // 'operator<' is defined for elem_type.
     void Sort() {
       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.
     // @param start  The index of the first element to destroy.
     // @param count  The number of elements to destroy.
     void DestructRange(index_type start, size_type count) {
       elem_type *iter = Elements() + start, *end = iter + count;
       for (; iter != end; ++iter) {
@@ -910,16 +981,46 @@ class nsTArray : public nsTArray_base {
     template<class Item>
     void AssignRange(index_type start, size_type count,
                      const Item *values) {
       elem_type *iter = Elements() + start, *end = iter + count;
       for (; iter != end; ++iter, ++values) {
         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>
 class nsAutoTArray : public nsTArray<E> {
   public:
     typedef nsTArray<E> base_type;
     typedef typename base_type::Header Header;
     typedef typename base_type::elem_type elem_type;
--- a/xpcom/tests/TestTArray.cpp
+++ b/xpcom/tests/TestTArray.cpp
@@ -518,16 +518,55 @@ static PRBool test_indexof() {
   array.AppendElement(5);
   array.RemoveElementAt(1);
   // we should not find the 5!
   return array.IndexOf(5, 1) == array.NoIndex;
 }
 
 //----
 
+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 }
 
 static const struct Test {
   const char* name;
   TestFunc    func;
 } tests[] = {
   DECL_TEST(test_int_array),
@@ -538,16 +577,17 @@ static const struct Test {
   DECL_TEST(test_string_array),
   DECL_TEST(test_comptr_array),
   DECL_TEST(test_refptr_array),
   DECL_TEST(test_ptrarray),
 #ifdef DEBUG
   DECL_TEST(test_autoarray),
 #endif
   DECL_TEST(test_indexof),
+  DECL_TEST(test_heap),
   { nsnull, nsnull }
 };
 
 }
 
 using namespace TestTArray;
 
 int main(int argc, char **argv) {