layout: Store most anonymous pseudo-elements in box slots (#37941)

Previously, anonymous boxes, such for anonymous table parts were not
associated with their non-pseudo ancestor DOM nodes. This presents a
problem when it comes time to clear layout data during incremental
layouts. This change reworks the way that pseudo-elements in general are
stored in their non-pseudo ancestor DOM nodes, allowing for any number
to be placed there.

This trades a bit of performance for space, as just adding a vector to
the node would add something like 24 bytes of storage to every node.
This change should have a neutral runtime memory usage.

Testing: This shouldn't change observable behavior and is thus covered
by
existing WPT tests. It will allow tests to pass in a subsequent PR.

Signed-off-by: Martin Robinson <mrobinson@igalia.com>
Co-authored-by: Oriol Brufau <obrufau@igalia.com>

Cargo.lock

@@ -4725,6 +4725,7 @@ dependencies = [
  "servo_geometry",
  "servo_malloc_size_of",
  "servo_url",
+ "smallvec",
  "stylo",
  "stylo_atoms",
  "stylo_traits",

components/layout/Cargo.toml

@@ -52,6 +52,7 @@ servo_arc = { workspace = true }
 servo_config = { path = "../config" }
 servo_geometry = { path = "../geometry" }
 servo_url = { path = "../url" }
+smallvec = { workspace = true }
 stylo = { workspace = true }
 stylo_atoms = { workspace = true }
 stylo_traits = { workspace = true }

components/layout/dom.rs

@@ -18,6 +18,7 @@ use malloc_size_of_derive::MallocSizeOf;
 use net_traits::image_cache::Image;
 use script::layout_dom::ServoLayoutNode;
 use servo_arc::Arc as ServoArc;
+use smallvec::SmallVec;
 use style::context::SharedStyleContext;
 use style::properties::ComputedValues;
 use style::selector_parser::{PseudoElement, RestyleDamage};
@@ -32,26 +33,35 @@ use crate::replaced::CanvasInfo;
 use crate::table::TableLevelBox;
 use crate::taffy::TaffyItemBox;
 
+#[derive(MallocSizeOf)]
+pub struct PseudoLayoutData {
+    pseudo: PseudoElement,
+    box_slot: ArcRefCell<Option<LayoutBox>>,
+}
+
 /// The data that is stored in each DOM node that is used by layout.
 #[derive(Default, MallocSizeOf)]
 pub struct InnerDOMLayoutData {
     pub(super) self_box: ArcRefCell<Option<LayoutBox>>,
-    pub(super) pseudo_before_box: ArcRefCell<Option<LayoutBox>>,
+    pub(super) pseudo_boxes: SmallVec<[PseudoLayoutData; 2]>,
-    pub(super) pseudo_after_box: ArcRefCell<Option<LayoutBox>>,
-    pub(super) pseudo_marker_box: ArcRefCell<Option<LayoutBox>>,
 }
 
 impl InnerDOMLayoutData {
     pub(crate) fn for_pseudo(
         &self,
         pseudo_element: Option<PseudoElement>,
-    ) -> AtomicRef<Option<LayoutBox>> {
+    ) -> Option<AtomicRef<Option<LayoutBox>>> {
-        match pseudo_element {
+        let Some(pseudo_element) = pseudo_element else {
-            Some(PseudoElement::Before) => self.pseudo_before_box.borrow(),
+            return Some(self.self_box.borrow());
-            Some(PseudoElement::After) => self.pseudo_after_box.borrow(),
+        };
-            Some(PseudoElement::Marker) => self.pseudo_marker_box.borrow(),
+
-            _ => self.self_box.borrow(),
+        for pseudo_layout_data in self.pseudo_boxes.iter() {
+            if pseudo_element == pseudo_layout_data.pseudo {
+                return Some(pseudo_layout_data.box_slot.borrow());
+            }
         }
+
+        None
     }
 }
 
@@ -171,16 +181,18 @@ pub struct BoxSlot<'dom> {
     pub(crate) marker: PhantomData<&'dom ()>,
 }
 
-/// A mutable reference to a `LayoutBox` stored in a DOM element.
+impl From<ArcRefCell<Option<LayoutBox>>> for BoxSlot<'_> {
-impl BoxSlot<'_> {
+    fn from(layout_box_slot: ArcRefCell<Option<LayoutBox>>) -> Self {
-    pub(crate) fn new(slot: ArcRefCell<Option<LayoutBox>>) -> Self {
+        let slot = Some(layout_box_slot);
-        let slot = Some(slot);
         Self {
             slot,
             marker: PhantomData,
         }
     }
+}
 
+/// A mutable reference to a `LayoutBox` stored in a DOM element.
+impl BoxSlot<'_> {
     pub(crate) fn dummy() -> Self {
         let slot = None;
         Self {
@@ -226,12 +238,14 @@ pub(crate) trait NodeExt<'dom> {
     fn layout_data_mut(&self) -> AtomicRefMut<'dom, InnerDOMLayoutData>;
     fn layout_data(&self) -> Option<AtomicRef<'dom, InnerDOMLayoutData>>;
     fn element_box_slot(&self) -> BoxSlot<'dom>;
-    fn pseudo_element_box_slot(&self, which: PseudoElement) -> BoxSlot<'dom>;
+    fn pseudo_element_box_slot(&self, pseudo_element: PseudoElement) -> BoxSlot<'dom>;
-    fn unset_pseudo_element_box(&self, which: PseudoElement);
 
-    /// Remove boxes for the element itself, and its `:before` and `:after` if any.
+    /// Remove boxes for the element itself, and all of its pseudo-element boxes.
     fn unset_all_boxes(&self);
 
+    /// Remove all pseudo-element boxes for this element.
+    fn unset_all_pseudo_boxes(&self);
+
     fn fragments_for_pseudo(&self, pseudo_element: Option<PseudoElement>) -> Vec<Fragment>;
     fn invalidate_cached_fragment(&self);
 
@@ -341,62 +355,55 @@ impl<'dom> NodeExt<'dom> for ServoLayoutNode<'dom> {
     }
 
     fn element_box_slot(&self) -> BoxSlot<'dom> {
-        BoxSlot::new(self.layout_data_mut().self_box.clone())
+        self.layout_data_mut().self_box.clone().into()
     }
 
-    fn pseudo_element_box_slot(&self, pseudo_element_type: PseudoElement) -> BoxSlot<'dom> {
+    fn pseudo_element_box_slot(&self, pseudo_element: PseudoElement) -> BoxSlot<'dom> {
-        let data = self.layout_data_mut();
+        let mut layout_data = self.layout_data_mut();
-        let cell = match pseudo_element_type {
+        let box_slot = ArcRefCell::new(None);
-            PseudoElement::Before => &data.pseudo_before_box,
+        layout_data.pseudo_boxes.push(PseudoLayoutData {
-            PseudoElement::After => &data.pseudo_after_box,
+            pseudo: pseudo_element,
-            PseudoElement::Marker => &data.pseudo_marker_box,
+            box_slot: box_slot.clone(),
-            _ => unreachable!(
+        });
-                "Asked for box slot for unsupported pseudo-element: {:?}",
+        box_slot.into()
-                pseudo_element_type
-            ),
-        };
-        BoxSlot::new(cell.clone())
-    }
-
-    fn unset_pseudo_element_box(&self, pseudo_element_type: PseudoElement) {
-        let data = self.layout_data_mut();
-        let cell = match pseudo_element_type {
-            PseudoElement::Before => &data.pseudo_before_box,
-            PseudoElement::After => &data.pseudo_after_box,
-            PseudoElement::Marker => &data.pseudo_marker_box,
-            _ => unreachable!(
-                "Asked for box slot for unsupported pseudo-element: {:?}",
-                pseudo_element_type
-            ),
-        };
-        *cell.borrow_mut() = None;
     }
 
     fn unset_all_boxes(&self) {
-        let data = self.layout_data_mut();
+        let mut layout_data = self.layout_data_mut();
-        *data.self_box.borrow_mut() = None;
+        *layout_data.self_box.borrow_mut() = None;
-        *data.pseudo_before_box.borrow_mut() = None;
+        layout_data.pseudo_boxes.clear();
-        *data.pseudo_after_box.borrow_mut() = None;
+
-        *data.pseudo_marker_box.borrow_mut() = None;
         // Stylo already takes care of removing all layout data
         // for DOM descendants of elements with `display: none`.
     }
 
+    fn unset_all_pseudo_boxes(&self) {
+        self.layout_data_mut().pseudo_boxes.clear();
+    }
+
     fn invalidate_cached_fragment(&self) {
         let data = self.layout_data_mut();
-        if let Some(data) = data.self_box.borrow_mut().as_mut() {
+        if let Some(data) = data.self_box.borrow_mut().as_ref() {
             data.invalidate_cached_fragment();
         }
+
+        for pseudo_layout_data in data.pseudo_boxes.iter() {
+            if let Some(layout_box) = pseudo_layout_data.box_slot.borrow().as_ref() {
+                layout_box.invalidate_cached_fragment();
+            }
+        }
     }
 
     fn fragments_for_pseudo(&self, pseudo_element: Option<PseudoElement>) -> Vec<Fragment> {
-        NodeExt::layout_data(self)
+        let Some(layout_data) = NodeExt::layout_data(self) else {
-            .and_then(|layout_data| {
+            return vec![];
-                layout_data
+        };
-                    .for_pseudo(pseudo_element)
+        let Some(layout_data) = layout_data.for_pseudo(pseudo_element) else {
-                    .as_ref()
+            return vec![];
-                    .map(LayoutBox::fragments)
+        };
-            })
+        layout_data
+            .as_ref()
+            .map(LayoutBox::fragments)
             .unwrap_or_default()
     }
 
@@ -407,21 +414,11 @@ impl<'dom> NodeExt<'dom> for ServoLayoutNode<'dom> {
             layout_object.repair_style(context, self, &style);
         }
 
-        if let Some(layout_object) = &*data.pseudo_before_box.borrow() {
+        for pseudo_layout_data in data.pseudo_boxes.iter() {
-            if let Some(node) = self.to_threadsafe().with_pseudo(PseudoElement::Before) {
+            if let Some(layout_box) = pseudo_layout_data.box_slot.borrow().as_ref() {
-                layout_object.repair_style(context, self, &node.style(context));
+                if let Some(node) = self.to_threadsafe().with_pseudo(pseudo_layout_data.pseudo) {
-            }
+                    layout_box.repair_style(context, self, &node.style(context));
-        }
+                }
-
-        if let Some(layout_object) = &*data.pseudo_after_box.borrow() {
-            if let Some(node) = self.to_threadsafe().with_pseudo(PseudoElement::After) {
-                layout_object.repair_style(context, self, &node.style(context));
-            }
-        }
-
-        if let Some(layout_object) = &*data.pseudo_marker_box.borrow() {
-            if let Some(node) = self.to_threadsafe().with_pseudo(PseudoElement::Marker) {
-                layout_object.repair_style(context, self, &node.style(context));
             }
         }
     }

components/layout/dom_traversal.rs

@@ -231,9 +231,7 @@ fn traverse_element<'dom>(
     context: &LayoutContext,
     handler: &mut impl TraversalHandler<'dom>,
 ) {
-    // Clear any existing pseudo-element box slot, because markers are not handled like
+    element.unset_all_pseudo_boxes();
-    // `::before`` and `::after`. They are processed during box tree creation.
-    element.unset_pseudo_element_box(PseudoElement::Marker);
 
     let replaced = ReplacedContents::for_element(element, context);
     let style = element.style(&context.style_context);
@@ -286,9 +284,6 @@ fn traverse_eager_pseudo_element<'dom>(
 ) {
     assert!(pseudo_element_type.is_eager());
 
-    // First clear any old contents from the node.
-    node_info.node.unset_pseudo_element_box(pseudo_element_type);
-
     // If this node doesn't have this eager pseudo-element, exit early. This depends on
     // the style applied to the element.
     let Some(pseudo_element_info) = node_info.pseudo(context, pseudo_element_type) else {
@@ -353,8 +348,8 @@ fn traverse_pseudo_element_contents<'dom>(
                     anonymous_info,
                     display_inline,
                     Contents::Replaced(contents),
-                    // We don’t keep pointers to boxes generated by contents of pseudo-elements
+                    info.node
-                    BoxSlot::dummy(),
+                        .pseudo_element_box_slot(PseudoElement::ServoAnonymousBox),
                 )
             },
         }

components/layout/flow/construct.rs

@@ -299,9 +299,12 @@ impl<'dom, 'style> BlockContainerBuilder<'dom, 'style> {
                 self.push_block_level_job_for_inline_formatting_context(inline_formatting_context);
             }
 
+            let box_slot = table_info
+                .node
+                .pseudo_element_box_slot(PseudoElement::ServoAnonymousTable);
             self.block_level_boxes.push(BlockLevelJob {
                 info: table_info,
-                box_slot: BoxSlot::dummy(),
+                box_slot,
                 kind: BlockLevelCreator::AnonymousTable { table_block },
                 propagated_data: self.propagated_data,
             });
@@ -683,10 +686,13 @@ impl<'dom> BlockContainerBuilder<'dom, '_> {
             })
             .clone();
 
+        let box_slot = self
+            .info
+            .node
+            .pseudo_element_box_slot(PseudoElement::ServoAnonymousBox);
         self.block_level_boxes.push(BlockLevelJob {
             info,
-            // FIXME(nox): We should be storing this somewhere.
+            box_slot,
-            box_slot: BoxSlot::dummy(),
             kind: BlockLevelCreator::SameFormattingContextBlock(
                 IntermediateBlockContainer::InlineFormattingContext(
                     BlockContainer::InlineFormattingContext(inline_formatting_context),

components/layout/flow/root.rs

@@ -302,10 +302,7 @@ impl<'dom> IncrementalBoxTreeUpdate<'dom> {
         }
 
         let layout_data = NodeExt::layout_data(&potential_dirty_root_node)?;
-        if layout_data.pseudo_before_box.borrow().is_some() {
+        if !layout_data.pseudo_boxes.is_empty() {
-            return None;
-        }
-        if layout_data.pseudo_after_box.borrow().is_some() {
             return None;
         }
 

components/layout/table/construct.rs

@@ -20,7 +20,7 @@ use super::{
 };
 use crate::cell::ArcRefCell;
 use crate::context::LayoutContext;
-use crate::dom::{BoxSlot, LayoutBox};
+use crate::dom::{BoxSlot, LayoutBox, NodeExt};
 use crate::dom_traversal::{Contents, NodeAndStyleInfo, NonReplacedContents, TraversalHandler};
 use crate::flow::{BlockContainerBuilder, BlockFormattingContext};
 use crate::formatting_contexts::{
@@ -715,12 +715,18 @@ impl<'style, 'dom> TableBuilderTraversal<'style, 'dom> {
         row_builder.finish();
 
         let style = anonymous_info.style.clone();
-        self.push_table_row(ArcRefCell::new(TableTrack {
+        let table_row = ArcRefCell::new(TableTrack {
             base: LayoutBoxBase::new((&anonymous_info).into(), style.clone()),
             group_index: self.current_row_group_index,
             is_anonymous: true,
             shared_background_style: SharedStyle::new(style),
-        }));
+        });
+        self.push_table_row(table_row.clone());
+
+        self.info
+            .node
+            .pseudo_element_box_slot(PseudoElement::ServoAnonymousTableRow)
+            .set(LayoutBox::TableLevelBox(TableLevelBox::Track(table_row)))
     }
 
     fn push_table_row(&mut self, table_track: ArcRefCell<TableTrack>) {
@@ -981,14 +987,22 @@ impl<'style, 'builder, 'dom, 'a> TableRowBuilder<'style, 'builder, 'dom, 'a> {
         }
 
         let block_container = builder.finish();
+        let new_table_cell = ArcRefCell::new(TableSlotCell {
+            base: LayoutBoxBase::new(BaseFragmentInfo::anonymous(), anonymous_info.style),
+            contents: BlockFormattingContext::from_block_container(block_container),
+            colspan: 1,
+            rowspan: 1,
+        });
         self.table_traversal
             .builder
-            .add_cell(ArcRefCell::new(TableSlotCell {
+            .add_cell(new_table_cell.clone());
-                base: LayoutBoxBase::new(BaseFragmentInfo::anonymous(), anonymous_info.style),
+
-                contents: BlockFormattingContext::from_block_container(block_container),
+        self.info
-                colspan: 1,
+            .node
-                rowspan: 1,
+            .pseudo_element_box_slot(PseudoElement::ServoAnonymousTableCell)
-            }));
+            .set(LayoutBox::TableLevelBox(TableLevelBox::Cell(
+                new_table_cell,
+            )));
     }
 }