Auto merge of #18484 - emilio:per-origin-sheets, r=SimonSapin

style: Store stylesheets per origin

This is the first step that will allow us to cache UA sheets across documents.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/servo/servo/18484)
<!-- Reviewable:end -->

components/script/dom/document.rs

@@ -141,7 +141,7 @@ use style::selector_parser::{RestyleDamage, Snapshot};
 use style::shared_lock::{SharedRwLock as StyleSharedRwLock, SharedRwLockReadGuard};
 use style::str::{HTML_SPACE_CHARACTERS, split_html_space_chars, str_join};
 use style::stylesheet_set::StylesheetSet;
-use style::stylesheets::{Stylesheet, StylesheetContents, OriginSet};
+use style::stylesheets::{Stylesheet, StylesheetContents, Origin, OriginSet};
 use task_source::TaskSource;
 use time;
 use timers::OneshotTimerCallback;
@@ -2435,7 +2435,7 @@ impl Document {
     pub fn stylesheet_at(&self, index: usize) -> Option<Root<CSSStyleSheet>> {
         let stylesheets = self.stylesheets.borrow();
 
-        stylesheets.get(index).and_then(|s| {
+        stylesheets.get(Origin::Author, index).and_then(|s| {
             s.owner.upcast::<Node>().get_cssom_stylesheet()
         })
     }

components/script/dom/htmllinkelement.rs

@@ -242,8 +242,8 @@ impl VirtualMethods for HTMLLinkElement {
             s.unbind_from_tree(context);
         }
 
-        if let Some(ref s) = *self.stylesheet.borrow() {
-            document_from_node(self).remove_stylesheet(self.upcast(), s);
+        if let Some(s) = self.stylesheet.borrow_mut().take() {
+            document_from_node(self).remove_stylesheet(self.upcast(), &s);
         }
     }
 }

components/script/dom/htmlmetaelement.rs

@@ -196,8 +196,8 @@ impl VirtualMethods for HTMLMetaElement {
         if context.tree_in_doc {
             self.process_referrer_attribute();
 
-            if let Some(ref s) = *self.stylesheet.borrow() {
-                document_from_node(self).remove_stylesheet(self.upcast(), s);
+            if let Some(s) = self.stylesheet.borrow_mut().take() {
+                document_from_node(self).remove_stylesheet(self.upcast(), &s);
             }
         }
     }

components/script/dom/htmlstyleelement.rs

@@ -190,9 +190,8 @@ impl VirtualMethods for HTMLStyleElement {
         }
 
         if context.tree_in_doc {
-            if let Some(ref s) = *self.stylesheet.borrow() {
-                let doc = document_from_node(self);
-                doc.remove_stylesheet(self.upcast(), s)
+            if let Some(s) = self.stylesheet.borrow_mut().take() {
+                document_from_node(self).remove_stylesheet(self.upcast(), &s)
             }
         }
     }

components/style/stylesheet_set.rs

@@ -9,7 +9,7 @@ use invalidation::stylesheets::StylesheetInvalidationSet;
 use media_queries::Device;
 use shared_lock::SharedRwLockReadGuard;
 use std::slice;
-use stylesheets::{Origin, OriginSet, PerOrigin, StylesheetInDocument};
+use stylesheets::{Origin, OriginSet, OriginSetIterator, PerOrigin, StylesheetInDocument};
 
 /// Entry for a StylesheetSet.
 #[cfg_attr(feature = "servo", derive(HeapSizeOf))]
@@ -32,10 +32,32 @@ where
 
 /// A iterator over the stylesheets of a list of entries in the StylesheetSet.
 #[derive(Clone)]
-pub struct StylesheetIterator<'a, S>(slice::Iter<'a, StylesheetSetEntry<S>>)
+pub struct StylesheetCollectionIterator<'a, S>(slice::Iter<'a, StylesheetSetEntry<S>>)
 where
     S: StylesheetInDocument + PartialEq + 'static;
 
+impl<'a, S> Iterator for StylesheetCollectionIterator<'a, S>
+where
+    S: StylesheetInDocument + PartialEq + 'static,
+{
+    type Item = &'a S;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.0.next().map(|entry| &entry.sheet)
+    }
+}
+
+/// An iterator over the flattened view of the stylesheet collections.
+#[derive(Clone)]
+pub struct StylesheetIterator<'a, S>
+where
+    S: StylesheetInDocument + PartialEq + 'static,
+{
+    origins: OriginSetIterator,
+    collections: &'a PerOrigin<SheetCollection<S>>,
+    current: Option<StylesheetCollectionIterator<'a, S>>,
+}
+
 impl<'a, S> Iterator for StylesheetIterator<'a, S>
 where
     S: StylesheetInDocument + PartialEq + 'static,
@@ -43,7 +65,23 @@ where
     type Item = &'a S;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.0.next().map(|entry| &entry.sheet)
+        loop {
+            if self.current.is_none() {
+                let next_origin = match self.origins.next() {
+                    Some(o) => o,
+                    None => return None,
+                };
+
+                self.current =
+                    Some(self.collections.borrow_for_origin(&next_origin).iter());
+            }
+
+            if let Some(s) = self.current.as_mut().unwrap().next() {
+                return Some(s)
+            }
+
+            self.current = None;
+        }
     }
 }
 
@@ -75,9 +113,12 @@ where
     'b: 'a,
     S: StylesheetInDocument + PartialEq + 'static,
 {
-    iter: slice::IterMut<'a, StylesheetSetEntry<S>>,
     guard: &'a SharedRwLockReadGuard<'b>,
-    origins_dirty: OriginSet,
+    origins_dirty: OriginSetIterator,
+    // NB: Bound to the StylesheetSet lifetime when constructed, see
+    // StylesheetSet::flush.
+    collections: *mut PerOrigin<SheetCollection<S>>,
+    current: Option<slice::IterMut<'a, StylesheetSetEntry<S>>>,
     origin_data_validity: PerOrigin<OriginValidity>,
     author_style_disabled: bool,
     had_invalidations: bool,
@@ -110,7 +151,7 @@ where
 
     /// Returns whether running the whole flushing process would be a no-op.
     pub fn nothing_to_do(&self) -> bool {
-        self.origins_dirty.is_empty()
+        self.origins_dirty.clone().next().is_none()
     }
 
     /// Returns whether any DOM invalidations were processed as a result of the
@@ -128,7 +169,7 @@ where
 {
     fn drop(&mut self) {
         debug_assert!(
-            self.iter.next().is_none(),
+            self.origins_dirty.next().is_none(),
             "You're supposed to fully consume the flusher"
         );
     }
@@ -145,9 +186,30 @@ where
         use std::mem;
 
         loop {
-            let potential_sheet = match self.iter.next() {
-                None => return None,
+            if self.current.is_none() {
+                let next_origin = match self.origins_dirty.next() {
+                    Some(o) => o,
+                    None => return None,
+                };
+
+                // Should've been cleared already.
+                debug_assert_eq!(
+                    unsafe { &*self.collections }
+                        .borrow_for_origin(&next_origin)
+                        .data_validity,
+                    OriginValidity::Valid
+                );
+
+                self.current =
+                    Some(unsafe { &mut *self.collections }.borrow_mut_for_origin(&next_origin).entries.iter_mut());
+            }
+
+            let potential_sheet = match self.current.as_mut().unwrap().next() {
                 Some(s) => s,
+                None => {
+                    self.current = None;
+                    continue;
+                }
             };
 
             let dirty = mem::replace(&mut potential_sheet.dirty, false);
@@ -158,10 +220,6 @@ where
             }
 
             let origin = potential_sheet.sheet.contents(self.guard).origin;
-            if !self.origins_dirty.contains(origin.into()) {
-                continue;
-            }
-
             if self.author_style_disabled && matches!(origin, Origin::Author) {
                 continue;
             }
@@ -177,32 +235,125 @@ where
     }
 }
 
-/// The set of stylesheets effective for a given document.
 #[cfg_attr(feature = "servo", derive(HeapSizeOf))]
-pub struct StylesheetSet<S>
+struct SheetCollection<S>
 where
     S: StylesheetInDocument + PartialEq + 'static,
 {
-    /// The actual list of all the stylesheets that apply to the given document,
-    /// each stylesheet associated with a unique ID.
+    /// The actual list of stylesheets.
     ///
     /// This is only a list of top-level stylesheets, and as such it doesn't
     /// include recursive `@import` rules.
     entries: Vec<StylesheetSetEntry<S>>,
 
-    /// The invalidations for stylesheets added or removed from this document.
-    invalidations: StylesheetInvalidationSet,
-
-    /// The origins whose stylesheets have changed so far.
-    origins_dirty: OriginSet,
-
     /// The validity of the data that was already there for a given origin.
     ///
     /// Note that an origin may appear on `origins_dirty`, but still have
     /// `OriginValidity::Valid`, if only sheets have been appended into it (in
     /// which case the existing data is valid, but the origin needs to be
     /// rebuilt).
-    origin_data_validity: PerOrigin<OriginValidity>,
+    data_validity: OriginValidity,
+}
+
+impl<S> Default for SheetCollection<S>
+where
+    S: StylesheetInDocument + PartialEq + 'static,
+{
+    fn default() -> Self {
+        Self {
+            entries: vec![],
+            data_validity: OriginValidity::Valid,
+        }
+    }
+}
+
+impl<S> SheetCollection<S>
+where
+    S: StylesheetInDocument + PartialEq + 'static,
+{
+    /// Returns the number of stylesheets in the set.
+    fn len(&self) -> usize {
+        self.entries.len()
+    }
+
+    /// Returns the `index`th stylesheet in the set if present.
+    fn get(&self, index: usize) -> Option<&S> {
+        self.entries.get(index).map(|e| &e.sheet)
+    }
+
+    fn remove(&mut self, sheet: &S) {
+        let old_len = self.entries.len();
+        self.entries.retain(|entry| entry.sheet != *sheet);
+        if cfg!(feature = "servo") {
+            // FIXME(emilio): Make Gecko's PresShell::AddUserSheet not suck.
+            //
+            // Hopefully that's not necessary for correctness, just somewhat
+            // overkill.
+            debug_assert!(self.entries.len() != old_len, "Sheet not found?");
+        }
+        // Removing sheets makes us tear down the whole cascade and invalidation
+        // data.
+        self.set_data_validity_at_least(OriginValidity::FullyInvalid);
+    }
+
+    fn contains(&self, sheet: &S) -> bool {
+        self.entries.iter().any(|e| e.sheet == *sheet)
+    }
+
+    /// Appends a given sheet into the collection.
+    fn append(&mut self, sheet: S) {
+        debug_assert!(!self.contains(&sheet));
+        self.entries.push(StylesheetSetEntry::new(sheet))
+        // Appending sheets doesn't alter the validity of the existing data, so
+        // we don't need to change `data_validity` here.
+    }
+
+    fn insert_before(&mut self, sheet: S, before_sheet: &S) {
+        debug_assert!(!self.contains(&sheet));
+
+        let index = self.entries.iter().position(|entry| {
+            entry.sheet == *before_sheet
+        }).expect("`before_sheet` stylesheet not found");
+
+        // Inserting stylesheets somewhere but at the end changes the validity
+        // of the cascade data, but not the invalidation data.
+        self.set_data_validity_at_least(OriginValidity::CascadeInvalid);
+        self.entries.insert(index, StylesheetSetEntry::new(sheet));
+    }
+
+    fn set_data_validity_at_least(&mut self, validity: OriginValidity) {
+        use std::cmp;
+        self.data_validity = cmp::max(validity, self.data_validity);
+    }
+
+    fn prepend(&mut self, sheet: S) {
+        debug_assert!(!self.contains(&sheet));
+        // Inserting stylesheets somewhere but at the end changes the validity
+        // of the cascade data, but not the invalidation data.
+        self.set_data_validity_at_least(OriginValidity::CascadeInvalid);
+        self.entries.insert(0, StylesheetSetEntry::new(sheet));
+    }
+
+    /// Returns an iterator over the current list of stylesheets.
+    fn iter(&self) -> StylesheetCollectionIterator<S> {
+        StylesheetCollectionIterator(self.entries.iter())
+    }
+}
+
+/// The set of stylesheets effective for a given document.
+#[cfg_attr(feature = "servo", derive(HeapSizeOf))]
+pub struct StylesheetSet<S>
+where
+    S: StylesheetInDocument + PartialEq + 'static,
+{
+    /// The collections of sheets per each origin.
+    collections: PerOrigin<SheetCollection<S>>,
+
+    /// The invalidations for stylesheets added or removed from this document.
+    invalidations: StylesheetInvalidationSet,
+
+    /// The origins whose stylesheets have changed so far.
+    origins_dirty: OriginSet,
 
     /// Has author style been disabled?
     author_style_disabled: bool,
@@ -215,22 +366,21 @@ where
     /// Create a new empty StylesheetSet.
     pub fn new() -> Self {
         StylesheetSet {
-            entries: vec![],
+            collections: Default::default(),
             invalidations: StylesheetInvalidationSet::new(),
             origins_dirty: OriginSet::empty(),
-            origin_data_validity: Default::default(),
             author_style_disabled: false,
         }
     }
 
     /// Returns the number of stylesheets in the set.
     pub fn len(&self) -> usize {
-        self.entries.len()
+        self.collections.iter_origins().fold(0, |s, (item, _)| s + item.len())
     }
 
-    /// Returns the number of stylesheets in the set.
-    pub fn get(&self, index: usize) -> Option<&S> {
-        self.entries.get(index).map(|s| &s.sheet)
+    /// Returns the `index`th stylesheet in the set for the given origin.
+    pub fn get(&self, origin: Origin, index: usize) -> Option<&S> {
+        self.collections.borrow_for_origin(&origin).get(index)
     }
 
     /// Returns whether author styles have been disabled for the current
@@ -239,10 +389,6 @@ where
         self.author_style_disabled
     }
 
-    fn remove_stylesheet_if_present(&mut self, sheet: &S) {
-        self.entries.retain(|entry| entry.sheet != *sheet);
-    }
-
     fn collect_invalidations_for(
         &mut self,
         device: Option<&Device>,
@@ -255,24 +401,6 @@ where
         self.origins_dirty |= sheet.contents(guard).origin;
     }
 
-    fn set_data_validity_at_least(
-        &mut self,
-        origin: Origin,
-        validity: OriginValidity,
-    ) {
-        use std::cmp;
-
-        debug_assert!(
-            self.origins_dirty.contains(origin.into()),
-            "data_validity should be a subset of origins_dirty"
-        );
-
-        let existing_validity =
-            self.origin_data_validity.borrow_mut_for_origin(&origin);
-
-        *existing_validity = cmp::max(*existing_validity, validity);
-    }
-
     /// Appends a new stylesheet to the current set.
     ///
     /// No device implies not computing invalidations.
@@ -283,11 +411,9 @@ where
         guard: &SharedRwLockReadGuard
     ) {
         debug!("StylesheetSet::append_stylesheet");
-        self.remove_stylesheet_if_present(&sheet);
         self.collect_invalidations_for(device, &sheet, guard);
-        // Appending sheets doesn't alter the validity of the existing data, so
-        // we don't need to change `origin_data_validity` here.
-        self.entries.push(StylesheetSetEntry::new(sheet));
+        let origin = sheet.contents(guard).origin;
+        self.collections.borrow_mut_for_origin(&origin).append(sheet);
     }
 
     /// Prepend a new stylesheet to the current set.
@@ -298,14 +424,10 @@ where
         guard: &SharedRwLockReadGuard
     ) {
         debug!("StylesheetSet::prepend_stylesheet");
-        self.remove_stylesheet_if_present(&sheet);
         self.collect_invalidations_for(device, &sheet, guard);
 
-        // Inserting stylesheets somewhere but at the end changes the validity
-        // of the cascade data, but not the invalidation data.
-        self.set_data_validity_at_least(sheet.contents(guard).origin, OriginValidity::CascadeInvalid);
-
-        self.entries.insert(0, StylesheetSetEntry::new(sheet));
+        let origin = sheet.contents(guard).origin;
+        self.collections.borrow_mut_for_origin(&origin).prepend(sheet)
     }
 
     /// Insert a given stylesheet before another stylesheet in the document.
@@ -317,16 +439,12 @@ where
         guard: &SharedRwLockReadGuard,
     ) {
         debug!("StylesheetSet::insert_stylesheet_before");
-        self.remove_stylesheet_if_present(&sheet);
-        let index = self.entries.iter().position(|entry| {
-            entry.sheet == before_sheet
-        }).expect("`before_sheet` stylesheet not found");
         self.collect_invalidations_for(device, &sheet, guard);
 
-        // Inserting stylesheets somewhere but at the end changes the validity
-        // of the cascade data, but not the invalidation data.
-        self.set_data_validity_at_least(sheet.contents(guard).origin, OriginValidity::CascadeInvalid);
-        self.entries.insert(index, StylesheetSetEntry::new(sheet));
+        let origin = sheet.contents(guard).origin;
+        self.collections
+            .borrow_mut_for_origin(&origin)
+            .insert_before(sheet, &before_sheet)
     }
 
     /// Remove a given stylesheet from the set.
@@ -337,13 +455,10 @@ where
         guard: &SharedRwLockReadGuard,
     ) {
         debug!("StylesheetSet::remove_stylesheet");
-        self.remove_stylesheet_if_present(&sheet);
-
         self.collect_invalidations_for(device, &sheet, guard);
 
-        // Removing sheets makes us tear down the whole cascade and invalidation
-        // data.
-        self.set_data_validity_at_least(sheet.contents(guard).origin, OriginValidity::FullyInvalid);
+        let origin = sheet.contents(guard).origin;
+        self.collections.borrow_mut_for_origin(&origin).remove(&sheet)
     }
 
     /// Notes that the author style has been disabled for this document.
@@ -377,17 +492,24 @@ where
         debug!("StylesheetSet::flush");
 
         let had_invalidations = self.invalidations.flush(document_element);
-        let origins_dirty = mem::replace(&mut self.origins_dirty, OriginSet::empty());
-        let origin_data_validity =
-            mem::replace(&mut self.origin_data_validity, Default::default());
+        let origins_dirty =
+            mem::replace(&mut self.origins_dirty, OriginSet::empty()).iter();
+
+        let mut origin_data_validity = PerOrigin::<OriginValidity>::default();
+        for origin in origins_dirty.clone() {
+            let collection = self.collections.borrow_mut_for_origin(&origin);
+            *origin_data_validity.borrow_mut_for_origin(&origin) =
+                mem::replace(&mut collection.data_validity, OriginValidity::Valid);
+        }
 
         StylesheetFlusher {
-            iter: self.entries.iter_mut(),
+            collections: &mut self.collections,
             author_style_disabled: self.author_style_disabled,
             had_invalidations,
             origins_dirty,
             origin_data_validity,
             guard,
+            current: None,
         }
     }
 
@@ -402,9 +524,13 @@ where
         mem::replace(&mut self.origins_dirty, OriginSet::empty())
     }
 
-    /// Returns an iterator over the current list of stylesheets.
+    /// Return an iterator over the flattened view of all the stylesheets.
     pub fn iter(&self) -> StylesheetIterator<S> {
-        StylesheetIterator(self.entries.iter())
+        StylesheetIterator {
+            origins: OriginSet::all().iter(),
+            collections: &self.collections,
+            current: None,
+        }
     }
 
     /// Mark the stylesheets for the specified origin as dirty, because
@@ -414,7 +540,9 @@ where
         self.origins_dirty |= origins;
         for origin in origins.iter() {
             // We don't know what happened, assume the worse.
-            self.set_data_validity_at_least(origin, OriginValidity::FullyInvalid);
+            self.collections
+                .borrow_mut_for_origin(&origin)
+                .set_data_validity_at_least(OriginValidity::FullyInvalid);
         }
     }
 }

components/style/stylesheets/mod.rs

@@ -42,7 +42,7 @@ pub use self::keyframes_rule::KeyframesRule;
 pub use self::loader::StylesheetLoader;
 pub use self::media_rule::MediaRule;
 pub use self::namespace_rule::NamespaceRule;
-pub use self::origin::{Origin, OriginSet, PerOrigin, PerOriginIter};
+pub use self::origin::{Origin, OriginSet, OriginSetIterator, PerOrigin, PerOriginIter};
 pub use self::page_rule::PageRule;
 pub use self::rule_parser::{State, TopLevelRuleParser};
 pub use self::rule_list::{CssRules, CssRulesHelpers};

components/style/stylesheets/origin.rs

@@ -78,6 +78,7 @@ impl BitOrAssign<Origin> for OriginSet {
 
 /// Iterates over the origins present in an `OriginSet`, in order from
 /// highest priority (author) to lower (user agent).
+#[derive(Clone)]
 pub struct OriginSetIterator {
     set: OriginSet,
     cur: i8,

components/style/stylist.rs

@@ -41,7 +41,7 @@ use smallvec::VecLike;
 use std::fmt::Debug;
 use std::ops;
 use style_traits::viewport::ViewportConstraints;
-use stylesheet_set::{OriginValidity, SheetRebuildKind, StylesheetSet, StylesheetIterator, StylesheetFlusher};
+use stylesheet_set::{OriginValidity, SheetRebuildKind, StylesheetSet, StylesheetFlusher};
 #[cfg(feature = "gecko")]
 use stylesheets::{CounterStyleRule, FontFaceRule, FontFeatureValuesRule, PageRule};
 use stylesheets::{CssRule, Origin, OriginSet, PerOrigin, PerOriginIter};
@@ -634,14 +634,6 @@ impl Stylist {
         self.stylesheets.force_dirty(origins)
     }
 
-    /// Iterate over the given set of stylesheets.
-    ///
-    /// This is very intentionally exposed only on `&mut self`, since we don't
-    /// want to give access to the stylesheet list from worker threads.
-    pub fn iter_stylesheets(&mut self) -> StylesheetIterator<StylistSheet> {
-        self.stylesheets.iter()
-    }
-
     /// Sets whether author style is enabled or not.
     pub fn set_author_style_disabled(&mut self, disabled: bool) {
         self.stylesheets.set_author_style_disabled(disabled);