Bug 1441136: Add a fast way to enumerate ShadowRoots in a document. r=smaug

MozReview-Commit-ID: 7QffP56jsyk

dom/base/ShadowRoot.cpp

@@ -81,12 +81,35 @@ ShadowRoot::~ShadowRoot()
     host->RemoveMutationObserver(this);
   }
 
+  if (IsComposedDocParticipant()) {
+    OwnerDoc()->RemoveComposedDocShadowRoot(*this);
+  }
+
+  MOZ_DIAGNOSTIC_ASSERT(!OwnerDoc()->IsComposedDocShadowRoot(*this));
+
   UnsetFlags(NODE_IS_IN_SHADOW_TREE);
 
   // nsINode destructor expects mSubtreeRoot == this.
   SetSubtreeRootPointer(this);
 }
 
+void
+ShadowRoot::SetIsComposedDocParticipant(bool aIsComposedDocParticipant)
+{
+  bool changed = mIsComposedDocParticipant != aIsComposedDocParticipant;
+  mIsComposedDocParticipant = aIsComposedDocParticipant;
+  if (!changed) {
+    return;
+  }
+
+  nsIDocument* doc = OwnerDoc();
+  if (IsComposedDocParticipant()) {
+    doc->AddComposedDocShadowRoot(*this);
+  } else {
+    doc->RemoveComposedDocShadowRoot(*this);
+  }
+}
+
 JSObject*
 ShadowRoot::WrapObject(JSContext* aCx, JS::Handle<JSObject*> aGivenProto)
 {

dom/base/ShadowRoot.h

@@ -186,10 +186,7 @@ public:
     return mIsComposedDocParticipant;
   }
 
-  void SetIsComposedDocParticipant(bool aIsComposedDocParticipant)
-  {
-    mIsComposedDocParticipant = aIsComposedDocParticipant;
-  }
+  void SetIsComposedDocParticipant(bool aIsComposedDocParticipant);
 
   nsresult GetEventTargetParent(EventChainPreVisitor& aVisitor) override;
 

dom/base/nsDocument.cpp

@@ -8638,6 +8638,7 @@ nsIDocument::DestroyElementMaps()
 #endif
   mStyledLinks.Clear();
   mIdentifierMap.Clear();
+  mComposedShadowRoots.Clear();
   mResponsiveContent.Clear();
   IncrementExpandoGeneration(*this);
 }

dom/base/nsIDocument.h

@@ -2967,6 +2967,33 @@ public:
     mResponsiveContent.RemoveEntry(aContent);
   }
 
+  void AddComposedDocShadowRoot(mozilla::dom::ShadowRoot& aShadowRoot)
+  {
+    MOZ_ASSERT(IsShadowDOMEnabled());
+    mComposedShadowRoots.PutEntry(&aShadowRoot);
+  }
+
+  using ShadowRootSet = nsTHashtable<nsPtrHashKey<mozilla::dom::ShadowRoot>>;
+
+  void RemoveComposedDocShadowRoot(mozilla::dom::ShadowRoot& aShadowRoot)
+  {
+    MOZ_ASSERT(IsShadowDOMEnabled());
+    mComposedShadowRoots.RemoveEntry(&aShadowRoot);
+  }
+
+  // If you're considering using this, you probably want to use
+  // ShadowRoot::IsComposedDocParticipant instead. This is just for
+  // sanity-checking.
+  bool IsComposedDocShadowRoot(mozilla::dom::ShadowRoot& aShadowRoot)
+  {
+    return mComposedShadowRoots.Contains(&aShadowRoot);
+  }
+
+  const ShadowRootSet& ComposedShadowRoots() const
+  {
+    return mComposedShadowRoots;
+  }
+
   // Notifies any responsive content added by AddResponsiveContent upon media
   // features values changing.
   void NotifyMediaFeatureValuesChanged();
@@ -3758,6 +3785,11 @@ protected:
   // Tracking for images in the document.
   RefPtr<mozilla::dom::ImageTracker> mImageTracker;
 
+  // A hashtable of ShadowRoots belonging to the composed doc.
+  //
+  // See ShadowRoot::SetIsComposedDocParticipant.
+  ShadowRootSet mComposedShadowRoots;
+
   // The set of all object, embed, video/audio elements or
   // nsIObjectLoadingContent or nsIDocumentActivity for which this is the owner
   // document. (They might not be in the document.)

layout/style/ServoStyleSet.cpp

@@ -168,26 +168,6 @@ ServoStyleSet::Init(nsPresContext* aPresContext)
   SetStylistXBLStyleSheetsDirty();
 }
 
-template<typename Functor>
-void
-EnumerateShadowRootsInSubtree(const nsINode& aRoot, const Functor& aCb)
-{
-  for (const nsINode* cur = &aRoot; cur; cur = cur->GetNextNode()) {
-    if (!cur->IsElement()) {
-      continue;
-    }
-
-    auto* shadowRoot = cur->AsElement()->GetShadowRoot();
-    if (!shadowRoot) {
-      continue;
-    }
-
-    aCb(*shadowRoot);
-    EnumerateShadowRootsInSubtree(*shadowRoot, aCb);
-  }
-}
-
-// FIXME(emilio): We may want a faster way to do this.
 template<typename Functor>
 void
 EnumerateShadowRoots(const nsIDocument& aDoc, const Functor& aCb)
@@ -196,7 +176,13 @@ EnumerateShadowRoots(const nsIDocument& aDoc, const Functor& aCb)
     return;
   }
 
-  EnumerateShadowRootsInSubtree(aDoc, aCb);
+  const nsIDocument::ShadowRootSet& shadowRoots = aDoc.ComposedShadowRoots();
+  for (auto iter = shadowRoots.ConstIter(); !iter.Done(); iter.Next()) {
+    ShadowRoot* root = iter.Get()->GetKey();
+    MOZ_ASSERT(root);
+    MOZ_DIAGNOSTIC_ASSERT(root->IsComposedDocParticipant());
+    aCb(*root);
+  }
 }
 
 void