Bug 1555867 - Inline one child in the rule tree. r=heycam

It is indeed the most common case according to a bit of measurement.

A non-atypical example from GitHub for example:

> Rule tree stats:
>  0 - 340
>  1 - 1403
>  2 - 28
>  3 - 8
>  4 - 2
>  6 - 1
>  7 - 3
>  8 - 2
>  12 - 2
>  14 - 1
>  41 - 1
>  45 - 1
>  67 - 1
>  68 - 1

Differential Revision: https://phabricator.services.mozilla.com/D33351

--HG--
extra : moz-landing-system : lando

servo/components/style/rule_tree/mod.rs

@@ -83,20 +83,16 @@ impl MallocSizeOf for RuleTree {
 
         while let Some(node) = stack.pop() {
             n += unsafe { ops.malloc_size_of(node.ptr()) };
-            stack.extend(unsafe {
-                (*node.ptr())
-                    .children
-                    .read()
-                    .iter()
-                    .map(|(_k, v)| v.clone())
-            });
+            unsafe {
+                (*node.ptr()).children.read().each(|c| stack.push(c.clone()));
+            }
         }
 
         n
     }
 }
 
-#[derive(Debug, Eq, Hash, PartialEq)]
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 struct ChildKey(CascadeLevel, ptr::NonNull<()>);
 
 unsafe impl Send for ChildKey {}
@@ -434,7 +430,7 @@ impl RuleTree {
         while let Some(node) = stack.pop() {
             let children = node.get().children.read();
             *children_count.entry(children.len()).or_insert(0) += 1;
-            stack.extend(children.iter().map(|(_k, v)| v.upgrade()))
+            children.each(|c| stack.push(c.upgrade()));
         }
 
         trace!("Rule tree stats:");
@@ -765,6 +761,149 @@ impl CascadeLevel {
     }
 }
 
+/// The children of a single rule node.
+///
+/// We optimize the case of no kids and a single child, since they're by far the
+/// most common case and it'd cause a bunch of bloat for no reason.
+///
+/// The children remove themselves when they go away, which means that it's ok
+/// for us to store weak pointers to them.
+enum RuleNodeChildren {
+    /// There are no kids.
+    Empty,
+    /// There's just one kid. This is an extremely common case, so we don't
+    /// bother allocating a map for it.
+    One(WeakRuleNode),
+    /// At least at one point in time there was more than one kid (that is to
+    /// say, we don't bother re-allocating if children are removed dynamically).
+    Map(Box<FxHashMap<ChildKey, WeakRuleNode>>),
+}
+
+impl Default for RuleNodeChildren {
+    fn default() -> Self {
+        RuleNodeChildren::Empty
+    }
+}
+
+impl RuleNodeChildren {
+    /// Executes a given function for each of the children.
+    fn each(&self, mut f: impl FnMut(&WeakRuleNode)) {
+        match *self {
+            RuleNodeChildren::Empty => {},
+            RuleNodeChildren::One(ref child) => f(child),
+            RuleNodeChildren::Map(ref map) => {
+                for (_key, kid) in map.iter() {
+                    f(kid)
+                }
+            },
+        }
+    }
+
+    fn len(&self) -> usize {
+        match *self {
+            RuleNodeChildren::Empty => 0,
+            RuleNodeChildren::One(..) => 1,
+            RuleNodeChildren::Map(ref map) => map.len(),
+        }
+    }
+
+    fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    fn get(&self, key: &ChildKey) -> Option<&WeakRuleNode> {
+        match *self {
+            RuleNodeChildren::Empty => return None,
+            RuleNodeChildren::One(ref kid) => {
+                // We're read-locked, so no need to do refcount stuff, since the
+                // child is only removed from the main thread, _and_ it'd need
+                // to write-lock us anyway.
+                if unsafe { (*kid.ptr()).key() } == *key {
+                    Some(kid)
+                } else {
+                    None
+                }
+            },
+            RuleNodeChildren::Map(ref map) => map.get(&key),
+        }
+    }
+
+    fn get_or_insert_with(
+        &mut self,
+        key: ChildKey,
+        get_new_child: impl FnOnce() -> StrongRuleNode,
+    ) -> StrongRuleNode {
+        let existing_child_key = match *self {
+            RuleNodeChildren::Empty => {
+                let new = get_new_child();
+                debug_assert_eq!(new.get().key(), key);
+                *self = RuleNodeChildren::One(new.downgrade());
+                return new;
+            },
+            RuleNodeChildren::One(ref weak) => unsafe {
+                // We're locked necessarily, so it's fine to look at our
+                // weak-child without refcount-traffic.
+                let existing_child_key = (*weak.ptr()).key();
+                if existing_child_key == key {
+                    return weak.upgrade();
+                }
+                existing_child_key
+            },
+            RuleNodeChildren::Map(ref mut map) => {
+                return match map.entry(key) {
+                    hash::map::Entry::Occupied(ref occupied) => occupied.get().upgrade(),
+                    hash::map::Entry::Vacant(vacant) => {
+                        let new = get_new_child();
+
+                        debug_assert_eq!(new.get().key(), key);
+                        vacant.insert(new.downgrade());
+
+                        new
+                    },
+                };
+            },
+        };
+
+        let existing_child = match mem::replace(self, RuleNodeChildren::Empty) {
+            RuleNodeChildren::One(o) => o,
+            _ => unreachable!(),
+        };
+        // Two rule-nodes are still a not-totally-uncommon thing, so
+        // avoid over-allocating entries.
+        //
+        // TODO(emilio): Maybe just inline two kids too?
+        let mut children = Box::new(FxHashMap::with_capacity_and_hasher(2, Default::default()));
+        children.insert(existing_child_key, existing_child);
+
+        let new = get_new_child();
+        debug_assert_eq!(new.get().key(), key);
+        children.insert(key, new.downgrade());
+
+        *self = RuleNodeChildren::Map(children);
+
+        new
+    }
+
+    fn remove(&mut self, key: &ChildKey) -> Option<WeakRuleNode> {
+        match *self {
+            RuleNodeChildren::Empty => return None,
+            RuleNodeChildren::One(ref one) => {
+                if unsafe { (*one.ptr()).key() } != *key {
+                    return None;
+                }
+            },
+            RuleNodeChildren::Map(ref mut multiple) => {
+                return multiple.remove(key);
+            },
+        }
+
+        match mem::replace(self, RuleNodeChildren::Empty) {
+            RuleNodeChildren::One(o) => Some(o),
+            _ => unreachable!(),
+        }
+    }
+}
+
 /// A node in the rule tree.
 pub struct RuleNode {
     /// The root node. Only the root has no root pointer, for obvious reasons.
@@ -790,7 +929,7 @@ pub struct RuleNode {
 
     /// The children of a given rule node. Children remove themselves from here
     /// when they go away.
-    children: RwLock<FxHashMap<ChildKey, WeakRuleNode>>,
+    children: RwLock<RuleNodeChildren>,
 
     /// The next item in the rule tree free list, that starts on the root node.
     ///
@@ -881,8 +1020,14 @@ impl RuleNode {
         }
     }
 
-    fn key(&self) -> Option<ChildKey> {
-        Some(ChildKey(self.level, self.source.as_ref()?.key()))
+    fn key(&self) -> ChildKey {
+        ChildKey(
+            self.level,
+            self.source
+                .as_ref()
+                .expect("Called key() on the root node")
+                .key(),
+        )
     }
 
     fn is_root(&self) -> bool {
@@ -906,7 +1051,7 @@ impl RuleNode {
         );
 
         if let Some(parent) = self.parent.as_ref() {
-            let weak = parent.get().children.write().remove(&self.key().unwrap());
+            let weak = parent.get().children.write().remove(&self.key());
             assert_eq!(weak.unwrap().ptr() as *const _, self as *const _);
         }
     }
@@ -943,12 +1088,12 @@ impl RuleNode {
         }
 
         let _ = write!(writer, "\n");
-        for (_, child) in self.children.read().iter() {
+        self.children.read().each(|child| {
             child
                 .upgrade()
                 .get()
                 .dump(guards, writer, indent + INDENT_INCREMENT);
-        }
+        });
     }
 }
 
@@ -1010,21 +1155,14 @@ impl StrongRuleNode {
             return child.upgrade();
         }
 
-        match RwLockUpgradableReadGuard::upgrade(read_guard).entry(key) {
-            hash::map::Entry::Occupied(ref occupied) => occupied.get().upgrade(),
-            hash::map::Entry::Vacant(vacant) => {
-                let new_node = StrongRuleNode::new(Box::new(RuleNode::new(
-                    root,
-                    self.clone(),
-                    source.clone(),
-                    level,
-                )));
-
-                vacant.insert(new_node.downgrade());
-
-                new_node
-            },
-        }
+        RwLockUpgradableReadGuard::upgrade(read_guard).get_or_insert_with(key, move || {
+            StrongRuleNode::new(Box::new(RuleNode::new(
+                root,
+                self.clone(),
+                source.clone(),
+                level,
+            )))
+        })
     }
 
     /// Raw pointer to the RuleNode

servo/ports/geckolib/tests/size_of.rs

@@ -53,7 +53,7 @@ size_of_test!(
     16
 );
 
-size_of_test!(test_size_of_rule_node, RuleNode, 88);
+size_of_test!(test_size_of_rule_node, RuleNode, 80);
 
 // This is huge, but we allocate it on the stack and then never move it,
 // we only pass `&mut SourcePropertyDeclaration` references around.