Bug 1426493 - Optimize Element.getElementById inside shadow trees. r=bzbarsky

This is the only optimization missing compared to documents.

Differential Revision: https://phabricator.services.mozilla.com/D12215

--HG--
extra : moz-landing-system : lando

dom/base/nsINode.cpp

@@ -2806,43 +2806,40 @@ nsINode::ParseSelectorList(const nsAString& aSelectorString,
   return ret;
 }
 
-namespace {
-struct SelectorMatchInfo {
-};
-} // namespace
-
 // Given an id, find first element with that id under aRoot.
 // If none found, return nullptr. aRoot must be in the document.
 inline static Element*
-FindMatchingElementWithId(const nsAString& aId, nsINode* aRoot)
+FindMatchingElementWithId(const nsAString& aId,
+                          const Element& aRoot,
+                          const DocumentOrShadowRoot& aContainingDocOrShadowRoot)
 {
-  MOZ_ASSERT(aRoot->IsInUncomposedDoc(),
-             "Don't call me if the root is not in the document");
-  // FIXME(emilio): It'd be nice to optimize this for shadow roots too.
-  MOZ_ASSERT(aRoot->IsElement() || aRoot->IsDocument(),
-             "The optimization below to check ContentIsDescendantOf only for "
-             "elements depends on aRoot being either an element or a "
-             "document if it's in the document.  Note that document fragments "
-             "can't be IsInUncomposedDoc(), so should never show up here.");
+  MOZ_ASSERT(aRoot.SubtreeRoot() == &aContainingDocOrShadowRoot.AsNode());
+  MOZ_ASSERT(aRoot.IsInUncomposedDoc() || aRoot.IsInShadowTree(),
+             "Don't call me if the root is not in the document or in a shadow tree");
 
-  const nsTArray<Element*>* elements = aRoot->OwnerDoc()->GetAllElementsForId(aId);
+  const nsTArray<Element*>* elements =
+    aContainingDocOrShadowRoot.GetAllElementsForId(aId);
   if (!elements) {
     // Nothing to do; we're done
     return nullptr;
   }
 
-  // XXXbz: Should we fall back to the tree walk if aRoot is not the
-  // document and |elements| is long, for some value of "long"?
-  for (size_t i = 0; i < elements->Length(); ++i) {
-    Element* element = (*elements)[i];
-    if (!aRoot->IsElement() ||
-        (element != aRoot &&
-           nsContentUtils::ContentIsDescendantOf(element, aRoot))) {
-      // We have an element with the right id and it's a strict descendant
-      // of aRoot.
-      return element;
+  // XXXbz: Should we fall back to the tree walk if |elements| is long,
+  // for some value of "long"?
+  for (Element* element : *elements) {
+    if (MOZ_UNLIKELY(element == &aRoot)) {
+      continue;
     }
+
+    if (!nsContentUtils::ContentIsDescendantOf(element, &aRoot)) {
+      continue;
+    }
+
+    // We have an element with the right id and it's a strict descendant
+    // of aRoot.
+    return element;
   }
+
   return nullptr;
 }
 
@@ -2881,10 +2878,17 @@ nsINode::QuerySelectorAll(const nsAString& aSelector, ErrorResult& aResult)
 Element*
 nsINode::GetElementById(const nsAString& aId)
 {
+  MOZ_ASSERT(!IsShadowRoot(), "Should use the faster version");
   MOZ_ASSERT(IsElement() || IsDocumentFragment(),
              "Bogus this object for GetElementById call");
   if (IsInUncomposedDoc()) {
-    return FindMatchingElementWithId(aId, this);
+    MOZ_ASSERT(IsElement(), "Huh? A fragment in a document?");
+    return FindMatchingElementWithId(aId, *AsElement(), *OwnerDoc());
+  }
+
+  if (ShadowRoot* containingShadow = AsContent()->GetContainingShadow()) {
+    MOZ_ASSERT(IsElement(), "Huh? A fragment in a ShadowRoot?");
+    return FindMatchingElementWithId(aId, *AsElement(), *containingShadow);
   }
 
   for (nsIContent* kid = GetFirstChild(); kid; kid = kid->GetNextNode(this)) {