Bug 668855, part 5: scan weak maps. r=gal

xpcom/base/nsCycleCollector.cpp

@@ -1092,6 +1092,7 @@ struct nsCycleCollector
     void SelectPurple(GCGraphBuilder &builder);
     void MarkRoots(GCGraphBuilder &builder);
     void ScanRoots();
+    void ScanWeakMaps();
 
     // returns whether anything was collected
     bool CollectWhite(nsICycleCollectorListener *aListener);
@@ -1989,6 +1990,38 @@ struct scanVisitor
     PRUint32 &mWhiteNodeCount;
 };
 
+// Iterate over the WeakMaps.  If we mark anything while iterating
+// over the WeakMaps, we must iterate over all of the WeakMaps again.
+void
+nsCycleCollector::ScanWeakMaps()
+{
+    bool anyChanged;
+    do {
+        anyChanged = false;
+        for (PRUint32 i = 0; i < mGraph.mWeakMaps.Length(); i++) {
+            WeakMapping *wm = &mGraph.mWeakMaps[i];
+
+            // If mMap or mKey are null, the original object was marked black.
+            uint32 mColor = wm->mMap ? wm->mMap->mColor : black;
+            uint32 kColor = wm->mKey ? wm->mKey->mColor : black;
+            PtrInfo *v = wm->mVal;
+
+            // All non-null weak mapping maps, keys and values are
+            // roots (in the sense of WalkFromRoots) in the cycle
+            // collector graph, and thus should have been colored
+            // either black or white in ScanRoots().
+            NS_ASSERTION(mColor != grey, "Uncolored weak map");
+            NS_ASSERTION(kColor != grey, "Uncolored weak map key");
+            NS_ASSERTION(v->mColor != grey, "Uncolored weak map value");
+
+            if (mColor == black && kColor == black && v->mColor != black) {
+                GraphWalker<ScanBlackVisitor>(ScanBlackVisitor(mWhiteNodeCount)).Walk(v);
+                anyChanged = true;
+            }
+        }
+    } while (anyChanged);
+}
+
 void
 nsCycleCollector::ScanRoots()
 {
@@ -1999,6 +2032,8 @@ nsCycleCollector::ScanRoots()
     // NodePool::Enumerator.
     GraphWalker<scanVisitor>(scanVisitor(mWhiteNodeCount)).WalkFromRoots(mGraph); 
 
+    ScanWeakMaps();
+
 #ifdef DEBUG_CC
     // Sanity check: scan should have colored all grey nodes black or
     // white. So we ensure we have no grey nodes at this point.