Pass pending restyles instead of draining them from layout

components/layout_thread/dom_wrapper.rs

@@ -50,7 +50,6 @@ use script::layout_exports::{
     LayoutDom, LayoutElementHelpers, LayoutNodeHelpers, LayoutShadowRootHelpers,
     RawLayoutElementHelpers,
 };
-use script_layout_interface::message::PendingRestyle;
 use script_layout_interface::wrapper_traits::{
     DangerousThreadSafeLayoutNode, GetLayoutData, LayoutNode,
 };
@@ -405,14 +404,6 @@ impl<'ld> ServoLayoutDocument<'ld> {
             .next()
     }
 
-    pub fn drain_pending_restyles(&self) -> Vec<(ServoLayoutElement<'ld>, PendingRestyle)> {
-        let elements = unsafe { self.document.drain_pending_restyles() };
-        elements
-            .into_iter()
-            .map(|(el, snapshot)| (ServoLayoutElement::from_layout_js(el), snapshot))
-            .collect()
-    }
-
     pub fn needs_paint_from_layout(&self) {
         unsafe { self.document.needs_paint_from_layout() }
     }

components/layout_thread/lib.rs

@@ -1450,17 +1450,19 @@ impl LayoutThread {
             guards.author.clone(),
         );
 
-        let restyles = document.drain_pending_restyles();
+        let restyles = std::mem::take(&mut data.pending_restyles);
         debug!("Draining restyles: {}", restyles.len());
 
         let mut map = SnapshotMap::new();
         let elements_with_snapshot: Vec<_> = restyles
             .iter()
             .filter(|r| r.1.snapshot.is_some())
-            .map(|r| r.0)
+            .map(|r| unsafe { ServoLayoutNode::new(&r.0).as_element().unwrap() })
             .collect();
 
         for (el, restyle) in restyles {
+            let el = unsafe { ServoLayoutNode::new(&el).as_element().unwrap() };
+
             // Propagate the descendant bit up the ancestors. Do this before
             // the restyle calculation so that we can also do it for new
             // unstyled nodes, which the descendants bit helps us find.

components/layout_thread_2020/dom_wrapper.rs

@@ -50,7 +50,6 @@ use script::layout_exports::{
     LayoutDom, LayoutElementHelpers, LayoutNodeHelpers, LayoutShadowRootHelpers,
     RawLayoutElementHelpers,
 };
-use script_layout_interface::message::PendingRestyle;
 use script_layout_interface::wrapper_traits::{
     DangerousThreadSafeLayoutNode, GetLayoutData, LayoutNode,
 };
@@ -412,14 +411,6 @@ impl<'ld> ServoLayoutDocument<'ld> {
             .next()
     }
 
-    pub fn drain_pending_restyles(&self) -> Vec<(ServoLayoutElement<'ld>, PendingRestyle)> {
-        let elements = unsafe { self.document.drain_pending_restyles() };
-        elements
-            .into_iter()
-            .map(|(el, snapshot)| (ServoLayoutElement::from_layout_js(el), snapshot))
-            .collect()
-    }
-
     pub fn needs_paint_from_layout(&self) {
         unsafe { self.document.needs_paint_from_layout() }
     }

components/layout_thread_2020/lib.rs

@@ -1096,16 +1096,19 @@ impl LayoutThread {
             guards.author.clone(),
         );
 
-        let restyles = document.drain_pending_restyles();
+        let restyles = std::mem::take(&mut data.pending_restyles);
+        debug!("Draining restyles: {}", restyles.len());
 
         let mut map = SnapshotMap::new();
         let elements_with_snapshot: Vec<_> = restyles
             .iter()
             .filter(|r| r.1.snapshot.is_some())
-            .map(|r| r.0)
+            .map(|r| unsafe { ServoLayoutNode::new(&r.0).as_element().unwrap() })
             .collect();
 
         for (el, restyle) in restyles {
+            let el = unsafe { ServoLayoutNode::new(&el).as_element().unwrap() };
+
             // Propagate the descendant bit up the ancestors. Do this before
             // the restyle calculation so that we can also do it for new
             // unstyled nodes, which the descendants bit helps us find.

components/script/dom/document.rs

@@ -71,7 +71,7 @@ use crate::dom::location::Location;
 use crate::dom::messageevent::MessageEvent;
 use crate::dom::mouseevent::MouseEvent;
 use crate::dom::node::{self, document_from_node, window_from_node, CloneChildrenFlag};
-use crate::dom::node::{LayoutNodeHelpers, Node, NodeDamage, NodeFlags, ShadowIncluding};
+use crate::dom::node::{Node, NodeDamage, NodeFlags, ShadowIncluding};
 use crate::dom::nodeiterator::NodeIterator;
 use crate::dom::nodelist::NodeList;
 use crate::dom::pagetransitionevent::PageTransitionEvent;
@@ -134,6 +134,7 @@ use profile_traits::ipc as profile_ipc;
 use profile_traits::time::{TimerMetadata, TimerMetadataFrameType, TimerMetadataReflowType};
 use ref_slice::ref_slice;
 use script_layout_interface::message::{Msg, PendingRestyle, ReflowGoal};
+use script_layout_interface::TrustedNodeAddress;
 use script_traits::{AnimationState, DocumentActivity, MouseButton, MouseEventType};
 use script_traits::{
     MsDuration, ScriptMsg, TouchEventType, TouchId, UntrustedNodeAddress, WheelDelta,
@@ -2606,7 +2607,6 @@ pub enum DocumentSource {
 #[allow(unsafe_code)]
 pub trait LayoutDocumentHelpers {
     unsafe fn is_html_document_for_layout(&self) -> bool;
-    unsafe fn drain_pending_restyles(&self) -> Vec<(LayoutDom<Element>, PendingRestyle)>;
     unsafe fn needs_paint_from_layout(&self);
     unsafe fn will_paint(&self);
     unsafe fn quirks_mode(&self) -> QuirksMode;
@@ -2623,22 +2623,6 @@ impl LayoutDocumentHelpers for LayoutDom<Document> {
         (*self.unsafe_get()).is_html_document
     }
 
-    #[inline]
-    #[allow(unrooted_must_root)]
-    unsafe fn drain_pending_restyles(&self) -> Vec<(LayoutDom<Element>, PendingRestyle)> {
-        let mut elements = (*self.unsafe_get())
-            .pending_restyles
-            .borrow_mut_for_layout();
-        // Elements were in a document when they were added to this list, but that
-        // may no longer be true when the next layout occurs.
-        let result = elements
-            .drain()
-            .map(|(k, v)| (k.to_layout(), v))
-            .filter(|&(ref k, _)| k.upcast::<Node>().get_flag(NodeFlags::IS_CONNECTED))
-            .collect();
-        result
-    }
-
     #[inline]
     unsafe fn needs_paint_from_layout(&self) {
         (*self.unsafe_get()).needs_paint.set(true)
@@ -3579,6 +3563,16 @@ impl Document {
             (None, None) => ElementLookupResult::None,
         }
     }
+
+    #[allow(unrooted_must_root)]
+    pub fn drain_pending_restyles(&self) -> Vec<(TrustedNodeAddress, PendingRestyle)> {
+        self.pending_restyles
+            .borrow_mut()
+            .drain()
+            .filter(|(k, _)| k.upcast::<Node>().get_flag(NodeFlags::IS_CONNECTED))
+            .map(|(k, v)| (k.upcast::<Node>().to_trusted_node_address(), v))
+            .collect()
+    }
 }
 
 impl Element {

components/script/dom/window.rs

@@ -1617,13 +1617,14 @@ impl Window {
             reflow_info: Reflow {
                 page_clip_rect: self.page_clip_rect.get(),
             },
-            document: self.Document().upcast::<Node>().to_trusted_node_address(),
+            document: document.upcast::<Node>().to_trusted_node_address(),
             stylesheets_changed,
             window_size: self.window_size.get(),
             origin: self.origin().immutable().clone(),
             reflow_goal,
             script_join_chan: join_chan,
-            dom_count: self.Document().dom_count(),
+            dom_count: document.dom_count(),
+            pending_restyles: document.drain_pending_restyles(),
         };
 
         self.layout_chan

components/script_layout_interface/message.rs

@@ -219,6 +219,8 @@ pub struct ScriptReflow {
     pub dom_count: u32,
     /// The current window origin
     pub origin: ImmutableOrigin,
+    /// Restyle snapshot map.
+    pub pending_restyles: Vec<(TrustedNodeAddress, PendingRestyle)>,
 }
 
 pub struct LayoutThreadInit {