style: Collect ::part() rules during CascadeData rebuilds.

Unlike for :host() or ::slotted(), or regular rules, we don't need a whole
SelectorMap<>, so gotta make the code a bit more generic.

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

components/selectors/parser.rs

@@ -606,6 +606,31 @@ impl<Impl: SelectorImpl> Selector<Impl> {
         self.0.header.header.is_part()
     }
 
+    #[inline]
+    pub fn part(&self) -> Option<&Impl::PartName> {
+        if !self.is_part() {
+            return None;
+        }
+
+        let mut iter = self.iter();
+        if self.has_pseudo_element() {
+            // Skip the pseudo-element.
+            for _ in &mut iter {}
+
+            let combinator = iter.next_sequence()?;
+            debug_assert_eq!(combinator, Combinator::PseudoElement);
+        }
+
+        for component in iter {
+            if let Component::Part(ref part) = *component {
+                return Some(part);
+            }
+        }
+
+        debug_assert!(false, "is_part() lied somehow?");
+        None
+    }
+
     #[inline]
     pub fn pseudo_element(&self) -> Option<&Impl::PseudoElement> {
         if !self.has_pseudo_element() {

components/style/stylist.rs

@@ -33,6 +33,7 @@ use crate::stylesheets::{CounterStyleRule, FontFaceRule, FontFeatureValuesRule,
 use crate::stylesheets::{CssRule, Origin, OriginSet, PerOrigin, PerOriginIter};
 use crate::thread_state::{self, ThreadState};
 use crate::{Atom, LocalName, Namespace, WeakAtom};
+use fallible::FallibleVec;
 use hashglobe::FailedAllocationError;
 #[cfg(feature = "gecko")]
 use malloc_size_of::MallocUnconditionalShallowSizeOf;
@@ -1641,9 +1642,9 @@ impl<'a> SelectorVisitor for StylistSelectorVisitor<'a> {
 
 /// A set of rules for element and pseudo-elements.
 #[derive(Debug, Default, MallocSizeOf)]
-struct ElementAndPseudoRules {
+struct GenericElementAndPseudoRules<Map> {
     /// Rules from stylesheets at this `CascadeData`'s origin.
-    element_map: SelectorMap<Rule>,
+    element_map: Map,
 
     /// Rules from stylesheets at this `CascadeData`'s origin that correspond
     /// to a given pseudo-element.
@@ -1651,39 +1652,30 @@ struct ElementAndPseudoRules {
     /// FIXME(emilio): There are a bunch of wasted entries here in practice.
     /// Figure out a good way to do a `PerNonAnonBox` and `PerAnonBox` (for
     /// `precomputed_values_for_pseudo`) without duplicating a lot of code.
-    pseudos_map: PerPseudoElementMap<Box<SelectorMap<Rule>>>,
+    pseudos_map: PerPseudoElementMap<Box<Map>>,
 }
 
-impl ElementAndPseudoRules {
+impl<Map: Default + MallocSizeOf> GenericElementAndPseudoRules<Map> {
     #[inline(always)]
-    fn insert(
+    fn for_insertion(&mut self, pseudo_element: Option<&PseudoElement>) -> &mut Map {
-        &mut self,
-        rule: Rule,
-        pseudo_element: Option<&PseudoElement>,
-        quirks_mode: QuirksMode,
-    ) -> Result<(), FailedAllocationError> {
         debug_assert!(
-            pseudo_element.map_or(true, |pseudo| !pseudo.is_precomputed() &&
+            pseudo_element.map_or(true, |pseudo| {
-                !pseudo.is_unknown_webkit_pseudo_element())
+                !pseudo.is_precomputed() && !pseudo.is_unknown_webkit_pseudo_element()
+            }),
+            "Precomputed pseudos should end up in precomputed_pseudo_element_decls, \
+             and unknown webkit pseudos should be discarded before getting here"
         );
 
-        let map = match pseudo_element {
+        match pseudo_element {
             None => &mut self.element_map,
             Some(pseudo) => self
                 .pseudos_map
-                .get_or_insert_with(pseudo, || Box::new(SelectorMap::new())),
+                .get_or_insert_with(pseudo, || Box::new(Default::default())),
-        };
-
-        map.insert(rule, quirks_mode)
         }
-
-    fn clear(&mut self) {
-        self.element_map.clear();
-        self.pseudos_map.clear();
     }
 
     #[inline]
-    fn rules(&self, pseudo: Option<&PseudoElement>) -> Option<&SelectorMap<Rule>> {
+    fn rules(&self, pseudo: Option<&PseudoElement>) -> Option<&Map> {
         match pseudo {
             Some(pseudo) => self.pseudos_map.get(pseudo).map(|p| &**p),
             None => Some(&self.element_map),
@@ -1703,6 +1695,26 @@ impl ElementAndPseudoRules {
     }
 }
 
+type ElementAndPseudoRules = GenericElementAndPseudoRules<SelectorMap<Rule>>;
+type PartMap = PrecomputedHashMap<Atom, SmallVec<[Rule; 1]>>;
+type PartElementAndPseudoRules = GenericElementAndPseudoRules<PartMap>;
+
+impl ElementAndPseudoRules {
+    // TODO(emilio): Should we retain storage of these?
+    fn clear(&mut self) {
+        self.element_map.clear();
+        self.pseudos_map.clear();
+    }
+}
+
+impl PartElementAndPseudoRules {
+    // TODO(emilio): Should we retain storage of these?
+    fn clear(&mut self) {
+        self.element_map.clear();
+        self.pseudos_map.clear();
+    }
+}
+
 /// Data resulting from performing the CSS cascade that is specific to a given
 /// origin.
 ///
@@ -1727,6 +1739,12 @@ pub struct CascadeData {
     /// containing style scopes starting from the closest assigned slot.
     slotted_rules: Option<Box<ElementAndPseudoRules>>,
 
+    /// The data coming from ::part() pseudo-element rules.
+    ///
+    /// We need to store them separately because an element needs to match
+    /// ::part() pseudo-element rules in different shadow roots.
+    part_rules: Option<Box<PartElementAndPseudoRules>>,
+
     /// The invalidation map for these rules.
     invalidation_map: InvalidationMap,
 
@@ -1786,6 +1804,7 @@ impl CascadeData {
             normal_rules: ElementAndPseudoRules::default(),
             host_rules: None,
             slotted_rules: None,
+            part_rules: None,
             invalidation_map: InvalidationMap::new(),
             attribute_dependencies: PrecomputedHashSet::default(),
             state_dependencies: ElementState::empty(),
@@ -1876,6 +1895,12 @@ impl CascadeData {
         self.slotted_rules.as_ref().and_then(|d| d.rules(pseudo))
     }
 
+    /// Returns the parts rule map for a given pseudo-element.
+    #[inline]
+    pub fn part_rules(&self, pseudo: Option<&PseudoElement>) -> Option<&PartMap> {
+        self.part_rules.as_ref().and_then(|d| d.rules(pseudo))
+    }
+
     /// Collects all the applicable media query results into `results`.
     ///
     /// This duplicates part of the logic in `add_stylesheet`, which is
@@ -2005,10 +2030,22 @@ impl CascadeData {
                             }
                         }
 
+                        // Part is special, since given it doesn't have any
+                        // selectors inside, it's not worth using a whole
+                        // SelectorMap for it.
+                        if let Some(part) = selector.part() {
+                            self.part_rules
+                                .get_or_insert_with(|| Box::new(Default::default()))
+                                .for_insertion(pseudo_element)
+                                .try_entry(part.clone())?
+                                .or_insert_with(SmallVec::new)
+                                .try_push(rule)?;
+                        } else {
                             // NOTE(emilio): It's fine to look at :host and then at
                             // ::slotted(..), since :host::slotted(..) could never
                             // possibly match, as <slot> is not a valid shadow host.
-                        let rules = if selector.is_featureless_host_selector_or_pseudo_element() {
+                            let rules =
+                                if selector.is_featureless_host_selector_or_pseudo_element() {
                                     self.host_rules
                                         .get_or_insert_with(|| Box::new(Default::default()))
                                 } else if selector.is_slotted() {
@@ -2016,9 +2053,10 @@ impl CascadeData {
                                         .get_or_insert_with(|| Box::new(Default::default()))
                                 } else {
                                     &mut self.normal_rules
-                        };
+                                }
-
+                                .for_insertion(pseudo_element);
-                        rules.insert(rule, pseudo_element, quirks_mode)?;
+                            rules.insert(rule, quirks_mode)?;
+                        }
                     }
                     self.rules_source_order += 1;
                 },
@@ -2184,6 +2222,9 @@ impl CascadeData {
         if let Some(ref mut slotted_rules) = self.slotted_rules {
             slotted_rules.clear();
         }
+        if let Some(ref mut part_rules) = self.part_rules {
+            part_rules.clear();
+        }
         if let Some(ref mut host_rules) = self.host_rules {
             host_rules.clear();
         }
@@ -2212,6 +2253,9 @@ impl CascadeData {
         if let Some(ref slotted_rules) = self.slotted_rules {
             slotted_rules.add_size_of(ops, sizes);
         }
+        if let Some(ref part_rules) = self.part_rules {
+            part_rules.add_size_of(ops, sizes);
+        }
         if let Some(ref host_rules) = self.host_rules {
             host_rules.add_size_of(ops, sizes);
         }