script: Cache the `<iframe>` list per-Document (#34702)

This change creates a new struct `IFrameCollection` that is used to
cache the list of `<iframe>`s in a `Document` as long as the
`Document`'s DOM has not changed. This prevent constantly iterating the
entire DOM during *update the rendering*, which runs up to 60 times per
second as well as for other operations.

Signed-off-by: Martin Robinson <mrobinson@igalia.com>

7 files changed, 225 insertions, 107 deletions

diff --git a/components/script/document_collection.rs b/components/script/document_collection.rs
index 0dc4ce87bb4..d9c28f6b7db 100644
--- a/components/script/document_collection.rs
+++ b/components/script/document_collection.rs
@@ -55,8 +55,12 @@ impl DocumentCollection {
         pipeline_id: PipelineId,
         browsing_context_id: BrowsingContextId,
     ) -> Option<DomRoot<HTMLIFrameElement>> {
-        self.find_document(pipeline_id)
-            .and_then(|doc| doc.find_iframe(browsing_context_id))
+        self.find_document(pipeline_id).and_then(|document| {
+            document
+                .iframes()
+                .get(browsing_context_id)
+                .map(|iframe| iframe.element.as_rooted())
+        })
     }
 
     pub fn iter(&self) -> DocumentsIter<'_> {
@@ -83,9 +87,6 @@ impl DocumentCollection {
     ///
     /// [update-the-rendering]: https://html.spec.whatwg.org/multipage/#update-the-rendering
     pub(crate) fn documents_in_order(&self) -> Vec<PipelineId> {
-        // TODO: This is a fairly expensive operation, because iterating iframes requires walking
-        // the *entire* DOM for a document. Instead this should be cached and marked as dirty when
-        // the DOM of a document changes or documents are added or removed from our set.
         DocumentTree::new(self).documents_in_order()
     }
 }
@@ -140,7 +141,8 @@ impl DocumentTree {
         let mut tree = DocumentTree::default();
         for (id, document) in documents.iter() {
             let children: Vec<PipelineId> = document
-                .iter_iframes()
+                .iframes()
+                .iter()
                 .filter_map(|iframe| iframe.pipeline_id())
                 .filter(|iframe_pipeline_id| documents.find_document(*iframe_pipeline_id).is_some())
                 .collect();
diff --git a/components/script/dom/document.rs b/components/script/dom/document.rs
index 5dfcedda5b9..577b21d7dc2 100644
--- a/components/script/dom/document.rs
+++ b/components/script/dom/document.rs
@@ -15,7 +15,6 @@ use std::sync::{LazyLock, Mutex};
 use std::time::{Duration, Instant};
 
 use base::cross_process_instant::CrossProcessInstant;
-use base::id::BrowsingContextId;
 use canvas_traits::webgl::{self, WebGLContextId, WebGLMsg};
 use chrono::Local;
 use content_security_policy::{self as csp, CspList};
@@ -185,6 +184,7 @@ use crate::dom::window::Window;
 use crate::dom::windowproxy::WindowProxy;
 use crate::dom::xpathevaluator::XPathEvaluator;
 use crate::fetch::FetchCanceller;
+use crate::iframe_collection::IFrameCollection;
 use crate::network_listener::{NetworkListener, PreInvoke};
 use crate::realms::{enter_realm, AlreadyInRealm, InRealm};
 use crate::script_runtime::{CanGc, CommonScriptMsg, ScriptThreadEventCategory};
@@ -294,6 +294,8 @@ pub struct Document {
     scripts: MutNullableDom<HTMLCollection>,
     anchors: MutNullableDom<HTMLCollection>,
     applets: MutNullableDom<HTMLCollection>,
+    /// Information about the `<iframes>` in this [`Document`].
+    iframes: RefCell<IFrameCollection>,
     /// Lock use for style attributes and author-origin stylesheet objects in this document.
     /// Can be acquired once for accessing many objects.
     #[no_trace]
@@ -2251,9 +2253,12 @@ impl Document {
         }
         // Step 9
         if !recursive_flag {
-            for iframe in self.iter_iframes() {
+            // `prompt_to_unload` might cause futher modifications to the DOM so collecting here prevents
+            // a double borrow if the `IFrameCollection` needs to be validated again.
+            let iframes: Vec<_> = self.iframes().iter().collect();
+            for iframe in &iframes {
                 // TODO: handle the case of cross origin iframes.
-                let document = document_from_node(&*iframe);
+                let document = document_from_node(&**iframe);
                 can_unload = document.prompt_to_unload(true, can_gc);
                 if !document.salvageable() {
                     self.salvageable.set(false);
@@ -2320,9 +2325,12 @@ impl Document {
 
         // Step 13
         if !recursive_flag {
-            for iframe in self.iter_iframes() {
+            // `unload` might cause futher modifications to the DOM so collecting here prevents
+            // a double borrow if the `IFrameCollection` needs to be validated again.
+            let iframes: Vec<_> = self.iframes().iter().collect();
+            for iframe in &iframes {
                 // TODO: handle the case of cross origin iframes.
-                let document = document_from_node(&*iframe);
+                let document = document_from_node(&**iframe);
                 document.unload(true, can_gc);
                 if !document.salvageable() {
                     self.salvageable.set(false);
@@ -2677,7 +2685,7 @@ impl Document {
         self.loader.borrow_mut().inhibit_events();
 
         // Step 1.
-        for iframe in self.iter_iframes() {
+        for iframe in self.iframes().iter() {
             if let Some(document) = iframe.GetContentDocument() {
                 // TODO: abort the active documents of every child browsing context.
                 document.abort(can_gc);
@@ -2726,20 +2734,18 @@ impl Document {
         self.current_parser.get()
     }
 
-    /// Iterate over all iframes in the document.
-    pub fn iter_iframes(&self) -> impl Iterator<Item = DomRoot<HTMLIFrameElement>> {
-        self.upcast::<Node>()
-            .traverse_preorder(ShadowIncluding::Yes)
-            .filter_map(DomRoot::downcast::<HTMLIFrameElement>)
+    /// A reference to the [`IFrameCollection`] of this [`Document`], holding information about
+    /// `<iframe>`s found within it.
+    pub(crate) fn iframes(&self) -> Ref<IFrameCollection> {
+        self.iframes.borrow_mut().validate(self);
+        self.iframes.borrow()
     }
 
-    /// Find an iframe element in the document.
-    pub fn find_iframe(
-        &self,
-        browsing_context_id: BrowsingContextId,
-    ) -> Option<DomRoot<HTMLIFrameElement>> {
-        self.iter_iframes()
-            .find(|node| node.browsing_context_id() == Some(browsing_context_id))
+    /// A mutable reference to the [`IFrameCollection`] of this [`Document`], holding information about
+    /// `<iframe>`s found within it.
+    pub(crate) fn iframes_mut(&self) -> RefMut<IFrameCollection> {
+        self.iframes.borrow_mut().validate(self);
+        self.iframes.borrow_mut()
     }
 
     pub fn get_dom_interactive(&self) -> Option<CrossProcessInstant> {
@@ -3266,6 +3272,7 @@ impl Document {
             scripts: Default::default(),
             anchors: Default::default(),
             applets: Default::default(),
+            iframes: Default::default(),
             style_shared_lock: {
                 /// Per-process shared lock for author-origin stylesheets
                 ///
diff --git a/components/script/dom/htmliframeelement.rs b/components/script/dom/htmliframeelement.rs
index b3cf46343b5..9b4ad049c6e 100644
--- a/components/script/dom/htmliframeelement.rs
+++ b/components/script/dom/htmliframeelement.rs
@@ -804,7 +804,7 @@ impl VirtualMethods for HTMLIFrameElement {
                 );
                 let exited_window = exited_document.window();
                 exited_window.discard_browsing_context();
-                for exited_iframe in exited_document.iter_iframes() {
+                for exited_iframe in exited_document.iframes().iter() {
                     debug!("Discarding nested browsing context");
                     exited_iframe.destroy_nested_browsing_context();
                 }
diff --git a/components/script/dom/window.rs b/components/script/dom/window.rs
index d9b2a86e855..c7678cce0ad 100644
--- a/components/script/dom/window.rs
+++ b/components/script/dom/window.rs
@@ -51,14 +51,14 @@ use profile_traits::ipc as ProfiledIpc;
 use profile_traits::mem::ProfilerChan as MemProfilerChan;
 use profile_traits::time::ProfilerChan as TimeProfilerChan;
 use script_layout_interface::{
-    combine_id_with_fragment_type, FragmentType, IFrameSizes, Layout, PendingImageState, QueryMsg,
-    Reflow, ReflowGoal, ReflowRequest, TrustedNodeAddress,
+    combine_id_with_fragment_type, FragmentType, Layout, PendingImageState, QueryMsg, Reflow,
+    ReflowGoal, ReflowRequest, TrustedNodeAddress,
 };
 use script_traits::webdriver_msg::{WebDriverJSError, WebDriverJSResult};
 use script_traits::{
-    ConstellationControlMsg, DocumentState, IFrameSizeMsg, LoadData, NavigationHistoryBehavior,
-    ScriptMsg, ScriptToConstellationChan, ScrollState, StructuredSerializedData, Theme,
-    TimerSchedulerMsg, WindowSizeData, WindowSizeType,
+    ConstellationControlMsg, DocumentState, LoadData, NavigationHistoryBehavior, ScriptMsg,
+    ScriptToConstellationChan, ScrollState, StructuredSerializedData, Theme, TimerSchedulerMsg,
+    WindowSizeData, WindowSizeType,
 };
 use selectors::attr::CaseSensitivity;
 use servo_arc::Arc as ServoArc;
@@ -151,7 +151,7 @@ use crate::script_runtime::{
     CanGc, CommonScriptMsg, JSContext, Runtime, ScriptChan, ScriptPort, ScriptThreadEventCategory,
 };
 use crate::script_thread::{
-    with_script_thread, ImageCacheMsg, MainThreadScriptChan, MainThreadScriptMsg, ScriptThread,
+    ImageCacheMsg, MainThreadScriptChan, MainThreadScriptMsg, ScriptThread,
     SendableMainThreadScriptChan,
 };
 use crate::task_manager::TaskManager;
@@ -377,17 +377,6 @@ pub struct Window {
 
     /// <https://dom.spec.whatwg.org/#window-current-event>
     current_event: DomRefCell<Option<Dom<Event>>>,
-
-    /// Sizes of the various `<iframes>` that we might have on this [`Window`].
-    /// This is used to:
-    ///  - Let same-`ScriptThread` `<iframe>`s know synchronously when their
-    ///    size has changed, ensuring that the next layout has the right size.
-    ///  - Send the proper size for the `<iframe>` during new-Pipeline creation
-    ///    when the `src` attribute changes.
-    ///  - Let the `Constellation` know about `BrowsingContext` (one per `<iframe>`)
-    ///    size changes when an `<iframe>` changes size during layout.
-    #[no_trace]
-    iframe_sizes: RefCell<IFrameSizes>,
 }
 
 impl Window {
@@ -997,8 +986,7 @@ impl WindowMethods<crate::DomTypeHolder> for Window {
 
     // https://html.spec.whatwg.org/multipage/#accessing-other-browsing-contexts
     fn Length(&self) -> u32 {
-        let doc = self.Document();
-        doc.iter_iframes().count() as u32
+        self.Document().iframes().iter().count() as u32
     }
 
     // https://html.spec.whatwg.org/multipage/#dom-parent
@@ -1447,7 +1435,8 @@ impl WindowMethods<crate::DomTypeHolder> for Window {
 
         // https://html.spec.whatwg.org/multipage/#document-tree-child-browsing-context-name-property-set
         let iframes: Vec<_> = document
-            .iter_iframes()
+            .iframes()
+            .iter()
             .filter(|iframe| {
                 if let Some(window) = iframe.GetContentWindow() {
                     return window.get_name() == name;
@@ -1967,7 +1956,13 @@ impl Window {
             }
         }
 
-        self.handle_new_iframe_sizes_after_layout(results.iframe_sizes);
+        let size_messages = self
+            .Document()
+            .iframes_mut()
+            .handle_new_iframe_sizes_after_layout(results.iframe_sizes, self.device_pixel_ratio());
+        if !size_messages.is_empty() {
+            self.send_to_constellation(ScriptMsg::IFrameSizes(size_messages));
+        }
 
         document.update_animations_post_reflow();
         self.update_constellation_epoch();
@@ -1975,58 +1970,6 @@ impl Window {
         true
     }
 
-    /// Update the recorded iframe sizes of the contents of layout. When these sizes change,
-    /// send a message to the `Constellation` informing it of the new sizes.
-    fn handle_new_iframe_sizes_after_layout(&self, new_iframe_sizes: IFrameSizes) {
-        let old_iframe_sizes = self.iframe_sizes.replace(new_iframe_sizes);
-        let new_iframe_sizes = self.iframe_sizes.borrow();
-        if new_iframe_sizes.is_empty() {
-            return;
-        }
-
-        // Batch resize message to any local `Pipeline`s now, rather than waiting for them
-        // to filter asynchronously through the `Constellation`. This allows the new value
-        // to be reflected immediately in layout.
-        let device_pixel_ratio = self.device_pixel_ratio();
-        with_script_thread(|script_thread| {
-            for iframe_size in new_iframe_sizes.values() {
-                script_thread.handle_resize_message(
-                    iframe_size.pipeline_id,
-                    WindowSizeData {
-                        initial_viewport: iframe_size.size,
-                        device_pixel_ratio,
-                    },
-                    // TODO: This might send an extra resize event. This can fixed by explicitly
-                    // supporting `<iframe>` creation when the size isn't known yet.
-                    WindowSizeType::Resize,
-                );
-            }
-        });
-        // Send asynchronous updates to `Constellation.`
-        let size_messages: Vec<_> = new_iframe_sizes
-            .iter()
-            .filter_map(|(browsing_context_id, size)| {
-                match old_iframe_sizes.get(browsing_context_id) {
-                    Some(old_size) if old_size.size == size.size => None,
-                    Some(..) => Some(IFrameSizeMsg {
-                        browsing_context_id: *browsing_context_id,
-                        size: size.size,
-                        type_: WindowSizeType::Resize,
-                    }),
-                    None => Some(IFrameSizeMsg {
-                        browsing_context_id: *browsing_context_id,
-                        size: size.size,
-                        type_: WindowSizeType::Initial,
-                    }),
-                }
-            })
-            .collect();
-
-        if !size_messages.is_empty() {
-            self.send_to_constellation(ScriptMsg::IFrameSizes(size_messages));
-        }
-    }
-
     /// Reflows the page if it's possible to do so and the page is dirty. Returns true if layout
     /// actually happened, false otherwise.
     ///
@@ -2327,10 +2270,10 @@ impl Window {
     ) -> Option<Size2D<f32, CSSPixel>> {
         // Reflow might fail, but do a best effort to return the right size.
         self.layout_reflow(QueryMsg::InnerWindowDimensionsQuery, can_gc);
-        self.iframe_sizes
-            .borrow()
-            .get(&browsing_context_id)
-            .map(|iframe_size| iframe_size.size)
+        self.Document()
+            .iframes()
+            .get(browsing_context_id)
+            .and_then(|iframe| iframe.size)
     }
 
     #[allow(unsafe_code)]
@@ -2865,7 +2808,6 @@ impl Window {
             layout_marker: DomRefCell::new(Rc::new(Cell::new(true))),
             current_event: DomRefCell::new(None),
             theme: Cell::new(PrefersColorScheme::Light),
-            iframe_sizes: RefCell::default(),
         });
 
         unsafe { WindowBinding::Wrap(JSContext::from_ptr(runtime.cx()), win) }
diff --git a/components/script/iframe_collection.rs b/components/script/iframe_collection.rs
new file mode 100644
index 00000000000..08e1b326ced
--- /dev/null
+++ b/components/script/iframe_collection.rs
@@ -0,0 +1,165 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
+
+use std::cell::Cell;
+use std::default::Default;
+
+use base::id::BrowsingContextId;
+use euclid::{Scale, Size2D};
+use fnv::FnvHashMap;
+use script_layout_interface::IFrameSizes;
+use script_traits::{IFrameSizeMsg, WindowSizeData, WindowSizeType};
+use style_traits::CSSPixel;
+use webrender_api::units::DevicePixel;
+
+use crate::dom::bindings::inheritance::Castable;
+use crate::dom::bindings::root::{Dom, DomRoot};
+use crate::dom::htmliframeelement::HTMLIFrameElement;
+use crate::dom::node::{Node, ShadowIncluding};
+use crate::dom::types::Document;
+use crate::script_thread::with_script_thread;
+
+#[derive(JSTraceable, MallocSizeOf)]
+#[crown::unrooted_must_root_lint::must_root]
+pub(crate) struct IFrame {
+    pub element: Dom<HTMLIFrameElement>,
+    #[no_trace]
+    pub size: Option<Size2D<f32, CSSPixel>>,
+}
+
+#[derive(Default, JSTraceable, MallocSizeOf)]
+#[crown::unrooted_must_root_lint::must_root]
+pub(crate) struct IFrameCollection {
+    /// The version of the [`Document`] that this collection refers to. When the versions
+    /// do not match, the collection will need to be rebuilt.
+    document_version: Cell<u64>,
+    /// The `<iframe>`s in the collection.
+    iframes: Vec<IFrame>,
+}
+
+impl IFrameCollection {
+    /// Validate that the collection is up-to-date with the given [`Document`]. If it isn't up-to-date
+    /// rebuild it.
+    pub(crate) fn validate(&mut self, document: &Document) {
+        // TODO: Whether the DOM has changed at all can lead to rebuilding this collection
+        // when it isn't necessary. A better signal might be if any `<iframe>` nodes have
+        // been connected or disconnected.
+        let document_node = DomRoot::from_ref(document.upcast::<Node>());
+        let document_version = document_node.inclusive_descendants_version();
+        if document_version == self.document_version.get() {
+            return;
+        }
+
+        // Preserve any old sizes, but only for `<iframe>`s that already have a
+        // BrowsingContextId and a set size.
+        let mut old_sizes: FnvHashMap<_, _> = self
+            .iframes
+            .iter()
+            .filter_map(
+                |iframe| match (iframe.element.browsing_context_id(), iframe.size) {
+                    (Some(browsing_context_id), Some(size)) => Some((browsing_context_id, size)),
+                    _ => None,
+                },
+            )
+            .collect();
+
+        self.iframes = document_node
+            .traverse_preorder(ShadowIncluding::Yes)
+            .filter_map(DomRoot::downcast::<HTMLIFrameElement>)
+            .map(|element| {
+                let size = element
+                    .browsing_context_id()
+                    .and_then(|browsing_context_id| old_sizes.remove(&browsing_context_id));
+                IFrame {
+                    element: element.as_traced(),
+                    size,
+                }
+            })
+            .collect();
+        self.document_version.set(document_version);
+    }
+
+    pub(crate) fn get(&self, browsing_context_id: BrowsingContextId) -> Option<&IFrame> {
+        self.iframes
+            .iter()
+            .find(|iframe| iframe.element.browsing_context_id() == Some(browsing_context_id))
+    }
+
+    pub(crate) fn get_mut(
+        &mut self,
+        browsing_context_id: BrowsingContextId,
+    ) -> Option<&mut IFrame> {
+        self.iframes
+            .iter_mut()
+            .find(|iframe| iframe.element.browsing_context_id() == Some(browsing_context_id))
+    }
+
+    /// Set the size of an `<iframe>` in the collection given its `BrowsingContextId` and
+    /// the new size. Returns the old size.
+    pub(crate) fn set_size(
+        &mut self,
+        browsing_context_id: BrowsingContextId,
+        new_size: Size2D<f32, CSSPixel>,
+    ) -> Option<Size2D<f32, CSSPixel>> {
+        self.get_mut(browsing_context_id)
+            .expect("Tried to set a size for an unknown <iframe>")
+            .size
+            .replace(new_size)
+    }
+
+    /// Update the recorded iframe sizes of the contents of layout. Return a
+    /// [`Vec<IFrameSizeMsg`] containing the messages to send to the `Constellation`. A
+    /// message is only sent when the size actually changes.
+    pub(crate) fn handle_new_iframe_sizes_after_layout(
+        &mut self,
+        new_iframe_sizes: IFrameSizes,
+        device_pixel_ratio: Scale<f32, CSSPixel, DevicePixel>,
+    ) -> Vec<IFrameSizeMsg> {
+        if new_iframe_sizes.is_empty() {
+            return vec![];
+        }
+
+        new_iframe_sizes
+            .into_iter()
+            .filter_map(|(browsing_context_id, size)| {
+                // Batch resize message to any local `Pipeline`s now, rather than waiting for them
+                // to filter asynchronously through the `Constellation`. This allows the new value
+                // to be reflected immediately in layout.
+                let new_size = size.size;
+                with_script_thread(|script_thread| {
+                    script_thread.handle_resize_message(
+                        size.pipeline_id,
+                        WindowSizeData {
+                            initial_viewport: new_size,
+                            device_pixel_ratio,
+                        },
+                        WindowSizeType::Resize,
+                    );
+                });
+
+                let old_size = self.set_size(browsing_context_id, new_size);
+                // The `Constellation` should be up-to-date even when the in-ScriptThread pipelines
+                // might not be.
+                if old_size == Some(size.size) {
+                    return None;
+                }
+
+                let size_type = match old_size {
+                    Some(_) => WindowSizeType::Resize,
+                    None => WindowSizeType::Initial,
+                };
+
+                Some(IFrameSizeMsg {
+                    browsing_context_id,
+                    size: new_size,
+                    type_: size_type,
+                })
+            })
+            .collect()
+    }
+
+    pub(crate) fn iter(&self) -> impl Iterator<Item = DomRoot<HTMLIFrameElement>> + use<'_> {
+        self.iframes.iter().map(|iframe| iframe.element.as_rooted())
+    }
+}
diff --git a/components/script/lib.rs b/components/script/lib.rs
index 80a56236a2a..2e181725d27 100644
--- a/components/script/lib.rs
+++ b/components/script/lib.rs
@@ -55,6 +55,8 @@ mod layout_image;
 
 #[warn(deprecated)]
 pub mod document_collection;
+#[warn(deprecated)]
+pub mod iframe_collection;
 pub mod layout_dom;
 #[warn(deprecated)]
 mod mem;
diff --git a/components/script/script_thread.rs b/components/script/script_thread.rs
index 97fa9c55af3..1ff221a29ab 100644
--- a/components/script/script_thread.rs
+++ b/components/script/script_thread.rs
@@ -2988,15 +2988,15 @@ impl ScriptThread {
         browsing_context_id: BrowsingContextId,
         can_gc: CanGc,
     ) {
-        let doc = self
+        let document = self
             .documents
             .borrow()
             .find_document(parent_pipeline_id)
             .unwrap();
-        let frame_element = doc.find_iframe(browsing_context_id);
 
-        if let Some(ref frame_element) = frame_element {
-            doc.request_focus(Some(frame_element.upcast()), FocusType::Parent, can_gc);
+        let iframes = document.iframes();
+        if let Some(iframe) = iframes.get(browsing_context_id) {
+            document.request_focus(Some(iframe.element.upcast()), FocusType::Parent, can_gc);
         }
     }