style: Optimize QuerySelector in shadow trees.

Pretty much the same setup we have for document.

We have the awkwardness of having to check containing shadow manually for
ShadowRoot because it's not available in TNode (and making it available added a
bit more complexity that wasn't worth it IMO).

Bug: 1464428
Reviewed-by: xidorn
MozReview-Commit-ID: CqOh0sLHf6o

components/style/dom.rs

@@ -135,14 +135,17 @@ pub trait TDocument: Sized + Copy + Clone {
     fn quirks_mode(&self) -> QuirksMode;
 
     /// Get a list of elements with a given ID in this document, sorted by
-    /// document position.
+    /// tree position.
     ///
     /// Can return an error to signal that this list is not available, or also
     /// return an empty slice.
-    fn elements_with_id(
+    fn elements_with_id<'a>(
         &self,
         _id: &Atom,
-    ) -> Result<&[<Self::ConcreteNode as TNode>::ConcreteElement], ()> {
+    ) -> Result<&'a [<Self::ConcreteNode as TNode>::ConcreteElement], ()>
+    where
+        Self: 'a,
+    {
         Err(())
     }
 }
@@ -342,6 +345,21 @@ pub trait TShadowRoot: Sized + Copy + Clone + PartialEq {
     fn style_data<'a>(&self) -> &'a CascadeData
     where
         Self: 'a;
+
+    /// Get a list of elements with a given ID in this shadow root, sorted by
+    /// tree position.
+    ///
+    /// Can return an error to signal that this list is not available, or also
+    /// return an empty slice.
+    fn elements_with_id<'a>(
+        &self,
+        _id: &Atom,
+    ) -> Result<&'a [<Self::ConcreteNode as TNode>::ConcreteElement], ()>
+    where
+        Self: 'a,
+    {
+        Err(())
+    }
 }
 
 /// The element trait, the main abstraction the style crate acts over.

components/style/dom_apis.rs

@@ -221,14 +221,31 @@ where
     }
 }
 
-/// Returns whether a given element is descendant of a given `root` node.
+/// Returns whether a given element connected to `root` is descendant of `root`.
 ///
 /// NOTE(emilio): if root == element, this returns false.
-fn element_is_descendant_of<E>(element: E, root: E::ConcreteNode) -> bool
+fn connected_element_is_descendant_of<E>(element: E, root: E::ConcreteNode) -> bool
 where
     E: TElement,
 {
-    if element.as_node().is_in_document() && root == root.owner_doc().as_node() {
+    // Optimize for when the root is a document or a shadow root and the element
+    // is connected to that root.
+    if root.as_document().is_some() {
+        debug_assert!(element.as_node().is_in_document(), "Not connected?");
+        debug_assert_eq!(
+            root,
+            root.owner_doc().as_node(),
+            "Where did this element come from?",
+        );
+        return true;
+    }
+
+    if root.as_shadow_root().is_some() {
+        debug_assert_eq!(
+            element.containing_shadow().unwrap().as_node(),
+            root,
+            "Not connected?"
+        );
         return true;
     }
 
@@ -244,28 +261,33 @@ where
 }
 
 /// Fast path for iterating over every element with a given id in the document
-/// that `root` is connected to.
-fn fast_connected_elements_with_id<'a, D>(
-    doc: &'a D,
-    root: D::ConcreteNode,
+/// or shadow root that `root` is connected to.
+fn fast_connected_elements_with_id<'a, N>(
+    root: N,
     id: &Atom,
     quirks_mode: QuirksMode,
-) -> Result<&'a [<D::ConcreteNode as TNode>::ConcreteElement], ()>
+) -> Result<&'a [N::ConcreteElement], ()>
 where
-    D: TDocument,
+    N: TNode + 'a,
 {
-    debug_assert_eq!(root.owner_doc().as_node(), doc.as_node());
-
     let case_sensitivity = quirks_mode.classes_and_ids_case_sensitivity();
     if case_sensitivity != CaseSensitivity::CaseSensitive {
         return Err(());
     }
 
-    if !root.is_in_document() {
-        return Err(());
+    if root.is_in_document() {
+        return root.owner_doc().elements_with_id(id);
     }
 
-    doc.elements_with_id(id)
+    if let Some(shadow) = root.as_shadow_root() {
+        return shadow.elements_with_id(id);
+    }
+
+    if let Some(shadow) = root.as_element().and_then(|e| e.containing_shadow()) {
+        return shadow.elements_with_id(id);
+    }
+
+    Err(())
 }
 
 /// Collects elements with a given id under `root`, that pass `filter`.
@@ -280,8 +302,7 @@ fn collect_elements_with_id<E, Q, F>(
     Q: SelectorQuery<E>,
     F: FnMut(E) -> bool,
 {
-    let doc = root.owner_doc();
-    let elements = match fast_connected_elements_with_id(&doc, root, id, quirks_mode) {
+    let elements = match fast_connected_elements_with_id(root, id, quirks_mode) {
         Ok(elements) => elements,
         Err(()) => {
             let case_sensitivity = quirks_mode.classes_and_ids_case_sensitivity();
@@ -297,7 +318,7 @@ fn collect_elements_with_id<E, Q, F>(
     for element in elements {
         // If the element is not an actual descendant of the root, even though
         // it's connected, we don't really care about it.
-        if !element_is_descendant_of(*element, root) {
+        if !connected_element_is_descendant_of(*element, root) {
             continue;
         }
 
@@ -405,9 +426,8 @@ where
                         return Ok(());
                     }
 
-                    let doc = root.owner_doc();
-                    let elements = fast_connected_elements_with_id(&doc, root, id, quirks_mode)?;
-
+                    let elements =
+                        fast_connected_elements_with_id(root, id, quirks_mode)?;
                     if elements.is_empty() {
                         return Ok(());
                     }
@@ -432,7 +452,7 @@ where
                         //
                         // Give up on trying to optimize based on this id and
                         // keep walking our selector.
-                        if !element_is_descendant_of(*element, root) {
+                        if !connected_element_is_descendant_of(*element, root) {
                             continue 'component_loop;
                         }
 

components/style/gecko/wrapper.rs

@@ -87,6 +87,29 @@ use std::ptr;
 use string_cache::{Atom, Namespace, WeakAtom, WeakNamespace};
 use stylist::CascadeData;
 
+
+#[inline]
+fn elements_with_id<'a, 'le>(
+    array: *const structs::nsTArray<*mut RawGeckoElement>,
+) -> &'a [GeckoElement<'le>] {
+    unsafe {
+        if array.is_null() {
+            return &[];
+        }
+
+        let elements: &[*mut RawGeckoElement] = &**array;
+
+        // NOTE(emilio): We rely on the in-memory representation of
+        // GeckoElement<'ld> and *mut RawGeckoElement being the same.
+        #[allow(dead_code)]
+        unsafe fn static_assert() {
+            mem::transmute::<*mut RawGeckoElement, GeckoElement<'static>>(0xbadc0de as *mut _);
+        }
+
+        mem::transmute(elements)
+    }
+}
+
 /// A simple wrapper over `nsIDocument`.
 #[derive(Clone, Copy)]
 pub struct GeckoDocument<'ld>(pub &'ld structs::nsIDocument);
@@ -109,24 +132,14 @@ impl<'ld> TDocument for GeckoDocument<'ld> {
         self.0.mCompatMode.into()
     }
 
-    fn elements_with_id(&self, id: &Atom) -> Result<&[GeckoElement<'ld>], ()> {
-        unsafe {
-            let array = bindings::Gecko_GetElementsWithId(self.0, id.as_ptr());
-            if array.is_null() {
-                return Ok(&[]);
-            }
-
-            let elements: &[*mut RawGeckoElement] = &**array;
-
-            // NOTE(emilio): We rely on the in-memory representation of
-            // GeckoElement<'ld> and *mut RawGeckoElement being the same.
-            #[allow(dead_code)]
-            unsafe fn static_assert() {
-                mem::transmute::<*mut RawGeckoElement, GeckoElement<'static>>(0xbadc0de as *mut _);
-            }
-
-            Ok(mem::transmute(elements))
-        }
+    #[inline]
+    fn elements_with_id<'a>(&self, id: &Atom) -> Result<&'a [GeckoElement<'ld>], ()>
+    where
+        Self: 'a,
+    {
+        Ok(elements_with_id(unsafe {
+            bindings::Gecko_Document_GetElementsWithId(self.0, id.as_ptr())
+        }))
     }
 }
 
@@ -176,6 +189,16 @@ impl<'lr> TShadowRoot for GeckoShadowRoot<'lr> {
 
         &author_styles.data
     }
+
+    #[inline]
+    fn elements_with_id<'a>(&self, id: &Atom) -> Result<&'a [GeckoElement<'lr>], ()>
+    where
+        Self: 'a,
+    {
+        Ok(elements_with_id(unsafe {
+            bindings::Gecko_ShadowRoot_GetElementsWithId(self.0, id.as_ptr())
+        }))
+    }
 }
 
 /// A simple wrapper over a non-null Gecko node (`nsINode`) pointer.