style: Refactor the author sheet cache to keep alive the relevant StylesheetContents

This prevents incorrectly reusing cached results when the contents go
away and new contents are allocated with the same address.

Note that these keep alive transitively everything else under them, so
all other medialist keys don't need this.

By making this a proper hashmap it should also improve cache lookup
times if the cache grows too big.

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

components/servo_arc/lib.rs

@@ -248,6 +248,14 @@ impl<T> Arc<T> {
         }
     }
 
+    /// Like from_raw, but returns an addrefed arc instead.
+    #[inline]
+    pub unsafe fn from_raw_addrefed(ptr: *const T) -> Self {
+        let arc = Self::from_raw(ptr);
+        mem::forget(arc.clone());
+        arc
+    }
+
     /// Create a new static Arc<T> (one that won't reference count the object)
     /// and place it in the allocation provided by the specified `alloc`
     /// function.

components/style/stylesheets/stylesheet.rs

@@ -65,6 +65,10 @@ pub struct StylesheetContents {
     pub source_map_url: RwLock<Option<String>>,
     /// This stylesheet's source URL.
     pub source_url: RwLock<Option<String>>,
+
+    /// We don't want to allow construction outside of this file, to guarantee
+    /// that all contents are created with Arc<>.
+    _forbid_construction: (),
 }
 
 impl StylesheetContents {
@@ -82,7 +86,7 @@ impl StylesheetContents {
         use_counters: Option<&UseCounters>,
         allow_import_rules: AllowImportRules,
         sanitization_data: Option<&mut SanitizationData>,
-    ) -> Self {
+    ) -> Arc<Self> {
         let namespaces = RwLock::new(Namespaces::default());
         let (rules, source_map_url, source_url) = Stylesheet::parse_rules(
             css,
@@ -99,7 +103,7 @@ impl StylesheetContents {
             sanitization_data,
         );
 
-        Self {
+        Arc::new(Self {
             rules: CssRules::new(rules, &shared_lock),
             origin,
             url_data: RwLock::new(url_data),
@@ -107,7 +111,8 @@ impl StylesheetContents {
             quirks_mode,
             source_map_url: RwLock::new(source_map_url),
             source_url: RwLock::new(source_url),
-        }
+            _forbid_construction: (),
+        })
     }
 
     /// Creates a new StylesheetContents with the specified pre-parsed rules,
@@ -126,9 +131,9 @@ impl StylesheetContents {
         origin: Origin,
         url_data: UrlExtraData,
         quirks_mode: QuirksMode,
-    ) -> Self {
+    ) -> Arc<Self> {
         debug_assert!(rules.is_static());
-        Self {
+        Arc::new(Self {
             rules,
             origin,
             url_data: RwLock::new(url_data),
@@ -136,7 +141,8 @@ impl StylesheetContents {
             quirks_mode,
             source_map_url: RwLock::new(None),
             source_url: RwLock::new(None),
-        }
+            _forbid_construction: (),
+        })
     }
 
     /// Returns a reference to the list of rules.
@@ -178,6 +184,7 @@ impl DeepCloneWithLock for StylesheetContents {
             namespaces: RwLock::new((*self.namespaces.read()).clone()),
             source_map_url: RwLock::new((*self.source_map_url.read()).clone()),
             source_url: RwLock::new((*self.source_url.read()).clone()),
+            _forbid_construction: (),
         }
     }
 }
@@ -186,7 +193,7 @@ impl DeepCloneWithLock for StylesheetContents {
 #[derive(Debug)]
 pub struct Stylesheet {
     /// The contents of this stylesheet.
-    pub contents: StylesheetContents,
+    pub contents: Arc<StylesheetContents>,
     /// The lock used for objects inside this stylesheet
     pub shared_lock: SharedRwLock,
     /// List of media associated with the Stylesheet.
@@ -587,9 +594,9 @@ impl Clone for Stylesheet {
         // Make a deep clone of the media, using the new lock.
         let media = self.media.read_with(&guard).clone();
         let media = Arc::new(lock.wrap(media));
-        let contents = self
+        let contents = Arc::new(self
             .contents
-            .deep_clone_with_lock(&lock, &guard, &DeepCloneParams);
+            .deep_clone_with_lock(&lock, &guard, &DeepCloneParams));
 
         Stylesheet {
             contents,

components/style/stylist.rs

@@ -11,7 +11,7 @@ use crate::element_state::{DocumentState, ElementState};
 #[cfg(feature = "gecko")]
 use crate::gecko_bindings::structs::{ServoStyleSetSizes, StyleRuleInclusion};
 use crate::invalidation::element::invalidation_map::InvalidationMap;
-use crate::invalidation::media_queries::EffectiveMediaQueryResults;
+use crate::invalidation::media_queries::{EffectiveMediaQueryResults, MediaListKey, ToMediaListKey};
 use crate::invalidation::stylesheets::RuleChangeKind;
 use crate::media_queries::Device;
 use crate::properties::{self, CascadeMode, ComputedValues};
@@ -26,8 +26,7 @@ use crate::stylesheet_set::{DataValidity, DocumentStylesheetSet, SheetRebuildKin
 use crate::stylesheet_set::{DocumentStylesheetFlusher, SheetCollectionFlusher};
 use crate::stylesheets::keyframes_rule::KeyframesAnimation;
 use crate::stylesheets::viewport_rule::{self, MaybeNew, ViewportRule};
-use crate::stylesheets::StyleRule;
+use crate::stylesheets::{StyleRule, StylesheetInDocument, StylesheetContents};
-use crate::stylesheets::StylesheetInDocument;
 #[cfg(feature = "gecko")]
 use crate::stylesheets::{CounterStyleRule, FontFaceRule, FontFeatureValuesRule, PageRule};
 use crate::stylesheets::{CssRule, Origin, OriginSet, PerOrigin, PerOriginIter};
@@ -50,7 +49,9 @@ use smallbitvec::SmallBitVec;
 use smallvec::SmallVec;
 use std::sync::Mutex;
 use std::{mem, ops};
+use std::hash::{Hash, Hasher};
 use style_traits::viewport::ViewportConstraints;
+use fxhash::FxHashMap;
 
 /// The type of the stylesheets that the stylist contains.
 #[cfg(feature = "servo")]
@@ -60,6 +61,37 @@ pub type StylistSheet = crate::stylesheets::DocumentStyleSheet;
 #[cfg(feature = "gecko")]
 pub type StylistSheet = crate::gecko::data::GeckoStyleSheet;
 
+#[derive(Debug, Clone)]
+struct StylesheetContentsPtr(Arc<StylesheetContents>);
+
+impl PartialEq for StylesheetContentsPtr {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.0, &other.0)
+    }
+}
+
+impl Eq for StylesheetContentsPtr {}
+
+impl Hash for StylesheetContentsPtr {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        let contents: &StylesheetContents = &*self.0;
+        (contents as *const StylesheetContents).hash(state)
+    }
+}
+
+type StyleSheetContentList = Vec<StylesheetContentsPtr>;
+
+/// A key in the cascade data cache.
+#[derive(Debug, Hash, Default, PartialEq, Eq)]
+struct CascadeDataCacheKey {
+    media_query_results: Vec<MediaListKey>,
+    contents: StyleSheetContentList,
+}
+
+unsafe impl Send for CascadeDataCacheKey {}
+unsafe impl Sync for CascadeDataCacheKey {}
+
 trait CascadeDataCacheEntry : Sized {
     /// Returns a reference to the cascade data.
     fn cascade_data(&self) -> &CascadeData;
@@ -80,7 +112,7 @@ trait CascadeDataCacheEntry : Sized {
 }
 
 struct CascadeDataCache<Entry> {
-    entries: Vec<Arc<Entry>>,
+    entries: FxHashMap<CascadeDataCacheKey, Arc<Entry>>,
 }
 
 impl<Entry> CascadeDataCache<Entry>
@@ -88,7 +120,7 @@ where
     Entry: CascadeDataCacheEntry,
 {
     fn new() -> Self {
-        Self { entries: vec![] }
+        Self { entries: Default::default() }
     }
 
     fn len(&self) -> usize {
@@ -110,51 +142,66 @@ where
     where
         S: StylesheetInDocument + PartialEq + 'static,
     {
+        use std::collections::hash_map::Entry as HashMapEntry;
         debug!("StyleSheetCache::lookup({})", self.len());
 
         if !collection.dirty() {
             return Ok(None);
         }
 
-        let mut key = EffectiveMediaQueryResults::new();
+        let mut key = CascadeDataCacheKey::default();
         for sheet in collection.sheets() {
-            CascadeData::collect_applicable_media_query_results_into(device, sheet, guard, &mut key)
+            CascadeData::collect_applicable_media_query_results_into(
+                device,
+                sheet,
+                guard,
+                &mut key.media_query_results,
+                &mut key.contents,
+            )
         }
 
-        for entry in &self.entries {
+        let new_entry;
-            if std::ptr::eq(&**entry, old_entry) {
+        match self.entries.entry(key) {
-                // Avoid reusing our old entry (this can happen if we get
+            HashMapEntry::Vacant(e) => {
-                // invalidated due to CSSOM mutations and our old stylesheet
+                debug!("> Picking the slow path (not in the cache)");
-                // contents were already unique, for example). This old entry
+                new_entry = Entry::rebuild(
-                // will be pruned from the cache with take_unused() afterwards.
-                continue;
-            }
-            if entry.cascade_data().effective_media_query_results != key {
-                continue;
-            }
-            if log_enabled!(log::Level::Debug) {
-                debug!("cache hit for:");
-                for sheet in collection.sheets() {
-                    debug!(" > {:?}", sheet);
-                }
-            }
-            // The line below ensures the "committed" bit is updated properly
-            // below.
-            collection.each(|_, _| true);
-            return Ok(Some(entry.clone()));
-        }
-
-        debug!("> Picking the slow path");
-
-        let new_entry = Entry::rebuild(
                     device,
                     quirks_mode,
                     collection,
                     guard,
                     old_entry,
                 )?;
+                e.insert(new_entry.clone());
+            }
+            HashMapEntry::Occupied(mut e) => {
+                // Avoid reusing our old entry (this can happen if we get
+                // invalidated due to CSSOM mutations and our old stylesheet
+                // contents were already unique, for example).
+                if !std::ptr::eq(&**e.get(), old_entry) {
+                    if log_enabled!(log::Level::Debug) {
+                        debug!("cache hit for:");
+                        for sheet in collection.sheets() {
+                            debug!(" > {:?}", sheet);
+                        }
+                    }
+                    // The line below ensures the "committed" bit is updated
+                    // properly.
+                    collection.each(|_, _| true);
+                    return Ok(Some(e.get().clone()));
+                }
+
+                debug!("> Picking the slow path due to same entry as old");
+                new_entry = Entry::rebuild(
+                    device,
+                    quirks_mode,
+                    collection,
+                    guard,
+                    old_entry,
+                )?;
+                e.insert(new_entry.clone());
+            }
+        }
 
-        self.entries.push(new_entry.clone());
         Ok(Some(new_entry))
     }
 
@@ -167,25 +214,27 @@ where
     /// cache to not deadlock.
     fn take_unused(&mut self) -> SmallVec<[Arc<Entry>; 3]> {
         let mut unused = SmallVec::new();
-        for i in (0..self.entries.len()).rev() {
+        self.entries.retain(|_key, value| {
             // is_unique() returns false for static references, but we never
             // have static references to UserAgentCascadeDatas.  If we did, it
             // may not make sense to put them in the cache in the first place.
-            if self.entries[i].is_unique() {
+            if !value.is_unique() {
-                unused.push(self.entries.remove(i));
+                return true;
-            }
             }
+            unused.push(value.clone());
+            false
+        });
         unused
     }
 
-    fn take_all(&mut self) -> Vec<Arc<Entry>> {
+    fn take_all(&mut self) -> FxHashMap<CascadeDataCacheKey, Arc<Entry>> {
-        mem::replace(&mut self.entries, Vec::new())
+        mem::take(&mut self.entries)
     }
 
     #[cfg(feature = "gecko")]
     fn add_size_of(&self, ops: &mut MallocSizeOfOps, sizes: &mut ServoStyleSetSizes) {
         sizes.mOther += self.entries.shallow_size_of(ops);
-        for arc in self.entries.iter() {
+        for (_key, arc) in self.entries.iter() {
             // These are primary Arc references that can be measured
             // unconditionally.
             sizes.mOther += arc.unconditional_shallow_size_of(ops);
@@ -2033,7 +2082,8 @@ impl CascadeData {
         device: &Device,
         stylesheet: &S,
         guard: &SharedRwLockReadGuard,
-        results: &mut EffectiveMediaQueryResults,
+        results: &mut Vec<MediaListKey>,
+        contents_list: &mut StyleSheetContentList,
     ) where
         S: StylesheetInDocument + 'static,
     {
@@ -2042,19 +2092,25 @@ impl CascadeData {
         }
 
         debug!(" + {:?}", stylesheet);
-        results.saw_effective(stylesheet.contents());
+        let contents = stylesheet.contents();
+        results.push(contents.to_media_list_key());
+
+        // Safety: StyleSheetContents are reference-counted with Arc.
+        contents_list.push(StylesheetContentsPtr(unsafe {
+            Arc::from_raw_addrefed(contents)
+        }));
 
         for rule in stylesheet.effective_rules(device, guard) {
             match *rule {
                 CssRule::Import(ref lock) => {
                     let import_rule = lock.read_with(guard);
                     debug!(" + {:?}", import_rule.stylesheet.media(guard));
-                    results.saw_effective(import_rule);
+                    results.push(import_rule.to_media_list_key());
                 },
                 CssRule::Media(ref lock) => {
                     let media_rule = lock.read_with(guard);
                     debug!(" + {:?}", media_rule.media_queries.read_with(guard));
-                    results.saw_effective(media_rule);
+                    results.push(media_rule.to_media_list_key());
                 },
                 _ => {},
             }