1 files changed, 37 insertions, 33 deletions

diff --git a/components/style/stylesheet_set.rs b/components/style/stylesheet_set.rs
index 9e858a4566c..e2937e06e75 100644
--- a/components/style/stylesheet_set.rs
+++ b/components/style/stylesheet_set.rs
@@ -184,7 +184,9 @@ pub struct SheetCollectionFlusher<'a, S>
 where
     S: StylesheetInDocument + PartialEq + 'static,
 {
-    iter: slice::IterMut<'a, StylesheetSetEntry<S>>,
+    // TODO: This can be made an iterator again once
+    // https://github.com/rust-lang/rust/pull/82771 lands on stable.
+    entries: &'a mut [StylesheetSetEntry<S>],
     validity: DataValidity,
     dirty: bool,
 }
@@ -204,32 +206,42 @@ where
     pub fn data_validity(&self) -> DataValidity {
         self.validity
     }
+
+    /// Returns an iterator over the remaining list of sheets to consume.
+    pub fn sheets<'b>(&'b self) -> impl Iterator<Item = &'b S> {
+        self.entries.iter().map(|entry| &entry.sheet)
+    }
 }
 
-impl<'a, S> Iterator for SheetCollectionFlusher<'a, S>
+impl<'a, S> SheetCollectionFlusher<'a, S>
 where
     S: StylesheetInDocument + PartialEq + 'static,
 {
-    type Item = (&'a S, SheetRebuildKind);
-
-    fn next(&mut self) -> Option<Self::Item> {
-        loop {
-            let potential_sheet = self.iter.next()?;
-
+    /// Iterates over all sheets and values that we have to invalidate.
+    ///
+    /// TODO(emilio): This would be nicer as an iterator but we can't do that
+    /// until https://github.com/rust-lang/rust/pull/82771 stabilizes.
+    ///
+    /// Since we don't have a good use-case for partial iteration, this does the
+    /// trick for now.
+    pub fn each(self, mut callback: impl FnMut(&S, SheetRebuildKind) -> bool) {
+        for potential_sheet in self.entries.iter_mut() {
             let committed = mem::replace(&mut potential_sheet.committed, true);
-            if !committed {
+            let rebuild_kind = if !committed {
                 // If the sheet was uncommitted, we need to do a full rebuild
                 // anyway.
-                return Some((&potential_sheet.sheet, SheetRebuildKind::Full));
-            }
-
-            let rebuild_kind = match self.validity {
-                DataValidity::Valid => continue,
-                DataValidity::CascadeInvalid => SheetRebuildKind::CascadeOnly,
-                DataValidity::FullyInvalid => SheetRebuildKind::Full,
+                SheetRebuildKind::Full
+            } else {
+                match self.validity {
+                    DataValidity::Valid => continue,
+                    DataValidity::CascadeInvalid => SheetRebuildKind::CascadeOnly,
+                    DataValidity::FullyInvalid => SheetRebuildKind::Full,
+                }
             };
 
-            return Some((&potential_sheet.sheet, rebuild_kind));
+            if !callback(&potential_sheet.sheet, rebuild_kind) {
+                return;
+            }
         }
     }
 }
@@ -357,7 +369,7 @@ where
         let validity = mem::replace(&mut self.data_validity, DataValidity::Valid);
 
         SheetCollectionFlusher {
-            iter: self.entries.iter_mut(),
+            entries: &mut self.entries,
             dirty,
             validity,
         }
@@ -408,7 +420,7 @@ macro_rules! sheet_set_methods {
         ) {
             debug!(concat!($set_name, "::append_stylesheet"));
             self.collect_invalidations_for(device, &sheet, guard);
-            let collection = self.collection_for(&sheet, guard);
+            let collection = self.collection_for(&sheet);
             collection.append(sheet);
         }
 
@@ -423,7 +435,7 @@ macro_rules! sheet_set_methods {
             debug!(concat!($set_name, "::insert_stylesheet_before"));
             self.collect_invalidations_for(device, &sheet, guard);
 
-            let collection = self.collection_for(&sheet, guard);
+            let collection = self.collection_for(&sheet);
             collection.insert_before(sheet, &before_sheet);
         }
 
@@ -437,7 +449,7 @@ macro_rules! sheet_set_methods {
             debug!(concat!($set_name, "::remove_stylesheet"));
             self.collect_invalidations_for(device, &sheet, guard);
 
-            let collection = self.collection_for(&sheet, guard);
+            let collection = self.collection_for(&sheet);
             collection.remove(&sheet)
         }
 
@@ -487,7 +499,7 @@ macro_rules! sheet_set_methods {
                 RuleChangeKind::StyleRuleDeclarations => DataValidity::FullyInvalid,
             };
 
-            let collection = self.collection_for(&sheet, guard);
+            let collection = self.collection_for(&sheet);
             collection.set_data_validity_at_least(validity);
         }
     };
@@ -505,12 +517,8 @@ where
         }
     }
 
-    fn collection_for(
-        &mut self,
-        sheet: &S,
-        guard: &SharedRwLockReadGuard,
-    ) -> &mut SheetCollection<S> {
-        let origin = sheet.origin(guard);
+    fn collection_for(&mut self, sheet: &S) -> &mut SheetCollection<S> {
+        let origin = sheet.contents().origin;
         self.collections.borrow_mut_for_origin(&origin)
     }
 
@@ -658,11 +666,7 @@ where
         self.collection.len()
     }
 
-    fn collection_for(
-        &mut self,
-        _sheet: &S,
-        _guard: &SharedRwLockReadGuard,
-    ) -> &mut SheetCollection<S> {
+    fn collection_for(&mut self, _sheet: &S) -> &mut SheetCollection<S> {
         &mut self.collection
     }