From: Mounir Lamouri <mounir.lamouri@gmail.com>
Date: Tue, 1 Feb 2011 11:17:39 +0100
Subject: Bug 624187 - r=smaug a=dveditz

---
 docshell/base/nsDocShellEnumerator.cpp |   59 ++++++++++++++------------------
 docshell/base/nsDocShellEnumerator.h   |   19 ++++++----
 2 files changed, 37 insertions(+), 41 deletions(-)

diff --git a/docshell/base/nsDocShellEnumerator.cpp b/docshell/base/nsDocShellEnumerator.cpp
index ad1074b..2866457 100644
--- a/docshell/base/nsDocShellEnumerator.cpp
+++ b/docshell/base/nsDocShellEnumerator.cpp
@@ -44,16 +44,15 @@
 
 nsDocShellEnumerator::nsDocShellEnumerator(PRInt32 inEnumerationDirection)
 : mRootItem(nsnull)
-, mItemArray(nsnull)
 , mCurIndex(0)
 , mDocShellType(nsIDocShellTreeItem::typeAll)
+, mArrayValid(PR_FALSE)
 , mEnumerationDirection(inEnumerationDirection)
 {
 }
 
 nsDocShellEnumerator::~nsDocShellEnumerator()
 {
-  delete mItemArray;
 }
 
 NS_IMPL_ISUPPORTS1(nsDocShellEnumerator, nsISimpleEnumerator)
@@ -68,18 +67,14 @@ NS_IMETHODIMP nsDocShellEnumerator::GetNext(nsISupports **outCurItem)
   nsresult rv = EnsureDocShellArray();
   if (NS_FAILED(rv)) return rv;
   
-  if (mCurIndex >= 0 && mCurIndex < mItemArray->Count())
-  {
-    nsIDocShellTreeItem* thisItem = reinterpret_cast<nsIDocShellTreeItem*>(mItemArray->ElementAt(mCurIndex));
-    rv = thisItem->QueryInterface(NS_GET_IID(nsISupports), (void **)outCurItem);
-    if (NS_FAILED(rv)) return rv;
-  }
-  else
+  if (mCurIndex >= mItemArray.Length()) {
     return NS_ERROR_FAILURE;
-  
-  mCurIndex ++;
-  
-  return NS_OK;
+  }
+
+  // post-increment is important here
+  nsCOMPtr<nsISupports> item = do_QueryReferent(mItemArray[mCurIndex++], &rv);
+  item.forget(outCurItem);
+  return rv;
 }
 
 /* boolean hasMoreElements (); */
@@ -91,21 +86,21 @@ NS_IMETHODIMP nsDocShellEnumerator::HasMoreElements(PRBool *outHasMore)
   nsresult rv = EnsureDocShellArray();
   if (NS_FAILED(rv)) return rv;
 
-  *outHasMore = (mCurIndex < mItemArray->Count());
+  *outHasMore = (mCurIndex < mItemArray.Length());
   return NS_OK;
 }
 
 nsresult nsDocShellEnumerator::GetEnumerationRootItem(nsIDocShellTreeItem * *aEnumerationRootItem)
 {
   NS_ENSURE_ARG_POINTER(aEnumerationRootItem);
-  *aEnumerationRootItem = mRootItem;
-  NS_IF_ADDREF(*aEnumerationRootItem);
+  nsCOMPtr<nsIDocShellTreeItem> item = do_QueryReferent(mRootItem);
+  item.forget(aEnumerationRootItem);
   return NS_OK;
 }
 
 nsresult nsDocShellEnumerator::SetEnumerationRootItem(nsIDocShellTreeItem * aEnumerationRootItem)
 {
-  mRootItem = aEnumerationRootItem;
+  mRootItem = do_GetWeakReference(aEnumerationRootItem);
   ClearState();
   return NS_OK;
 }
@@ -132,12 +127,10 @@ nsresult nsDocShellEnumerator::First()
 
 nsresult nsDocShellEnumerator::EnsureDocShellArray()
 {
-  if (!mItemArray)
+  if (!mArrayValid)
   {
-    mItemArray = new nsVoidArray;
-    if (!mItemArray) return NS_ERROR_OUT_OF_MEMORY;
-  
-    return BuildDocShellArray(*mItemArray);
+    mArrayValid = PR_TRUE;
+    return BuildDocShellArray(mItemArray);
   }
   
   return NS_OK;
@@ -145,21 +138,21 @@ nsresult nsDocShellEnumerator::EnsureDocShellArray()
 
 nsresult nsDocShellEnumerator::ClearState()
 {
-  delete mItemArray;
-  mItemArray = nsnull;
-  
+  mItemArray.Clear();
+  mArrayValid = PR_FALSE;
   mCurIndex = 0;
   return NS_OK;
 }
 
-nsresult nsDocShellEnumerator::BuildDocShellArray(nsVoidArray& inItemArray)
+nsresult nsDocShellEnumerator::BuildDocShellArray(nsTArray<nsWeakPtr>& inItemArray)
 {
   NS_ENSURE_TRUE(mRootItem, NS_ERROR_NOT_INITIALIZED);
   inItemArray.Clear();
-  return BuildArrayRecursive(mRootItem, inItemArray);
+  nsCOMPtr<nsIDocShellTreeItem> item = do_QueryReferent(mRootItem);
+  return BuildArrayRecursive(item, inItemArray);
 }
 
-nsresult nsDocShellForwardsEnumerator::BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsVoidArray& inItemArray)
+nsresult nsDocShellForwardsEnumerator::BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsTArray<nsWeakPtr>& inItemArray)
 {
   nsresult rv;
   nsCOMPtr<nsIDocShellTreeNode> itemAsNode = do_QueryInterface(inItem, &rv);
@@ -170,8 +163,8 @@ nsresult nsDocShellForwardsEnumerator::BuildArrayRecursive(nsIDocShellTreeItem*
   if ((mDocShellType == nsIDocShellTreeItem::typeAll) ||
       (NS_SUCCEEDED(inItem->GetItemType(&itemType)) && (itemType == mDocShellType)))
   {
-    rv = inItemArray.AppendElement((void *)inItem);
-    if (NS_FAILED(rv)) return rv;
+    if (!inItemArray.AppendElement(do_GetWeakReference(inItem)))
+      return NS_ERROR_OUT_OF_MEMORY;
   }
 
   PRInt32   numChildren;
@@ -192,7 +185,7 @@ nsresult nsDocShellForwardsEnumerator::BuildArrayRecursive(nsIDocShellTreeItem*
 }
 
 
-nsresult nsDocShellBackwardsEnumerator::BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsVoidArray& inItemArray)
+nsresult nsDocShellBackwardsEnumerator::BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsTArray<nsWeakPtr>& inItemArray)
 {
   nsresult rv;
   nsCOMPtr<nsIDocShellTreeNode> itemAsNode = do_QueryInterface(inItem, &rv);
@@ -217,8 +210,8 @@ nsresult nsDocShellBackwardsEnumerator::BuildArrayRecursive(nsIDocShellTreeItem*
   if ((mDocShellType == nsIDocShellTreeItem::typeAll) ||
       (NS_SUCCEEDED(inItem->GetItemType(&itemType)) && (itemType == mDocShellType)))
   {
-    rv = inItemArray.AppendElement((void *)inItem);
-    if (NS_FAILED(rv)) return rv;
+    if (!inItemArray.AppendElement(do_GetWeakReference(inItem)))
+      return NS_ERROR_OUT_OF_MEMORY;
   }
 
 
diff --git a/docshell/base/nsDocShellEnumerator.h b/docshell/base/nsDocShellEnumerator.h
index aa4758a..eba2213 100644
--- a/docshell/base/nsDocShellEnumerator.h
+++ b/docshell/base/nsDocShellEnumerator.h
@@ -42,7 +42,9 @@
 #include "nsIEnumerator.h"
 
 #include "nsCOMPtr.h"
-#include "nsVoidArray.h"
+#include "nsTArray.h"
+#include "nsIWeakReference.h"
+#include "nsIWeakReferenceUtils.h"
 
 class nsIDocShellTreeItem;
 
@@ -98,17 +100,18 @@ protected:
   nsresult                    EnsureDocShellArray();
   nsresult                    ClearState();
   
-  nsresult                    BuildDocShellArray(nsVoidArray& inItemArray);
-  virtual nsresult            BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsVoidArray& inItemArray) = 0;
+  nsresult                    BuildDocShellArray(nsTArray<nsWeakPtr>& inItemArray);
+  virtual nsresult            BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsTArray<nsWeakPtr>& inItemArray) = 0;
     
 protected:
 
-  nsIDocShellTreeItem*        mRootItem;      // weak ref!
+  nsWeakPtr                   mRootItem;      // weak ref!
   
-  nsVoidArray*                mItemArray;     // flattened list of items with matching type
-  PRInt32                     mCurIndex;
+  nsTArray<nsWeakPtr>         mItemArray;     // flattened list of items with matching type
+  PRUint32                    mCurIndex;
   
   PRInt32                     mDocShellType;  // only want shells of this type
+  PRPackedBool                mArrayValid;    // is mItemArray up to date?
 
   const PRInt8                mEnumerationDirection;
 };
@@ -125,7 +128,7 @@ public:
 
 protected:
 
-  virtual nsresult            BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsVoidArray& inItemArray);
+  virtual nsresult            BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsTArray<nsWeakPtr>& inItemArray);
 
 };
 
@@ -139,6 +142,6 @@ public:
                               }
 protected:
 
-  virtual nsresult            BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsVoidArray& inItemArray);
+  virtual nsresult            BuildArrayRecursive(nsIDocShellTreeItem* inItem, nsTArray<nsWeakPtr>& inItemArray);
 
 };
