Auto merge of #18384 - emilio:invalidation-map-bloat, r=bholley

style: Don't waste a whole selector map for each class / id in the document.

On top of #18375, only last commit needs review.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/servo/servo/18384)
<!-- Reviewable:end -->

components/style/gecko/data.rs

@@ -16,7 +16,8 @@ use media_queries::{Device, MediaList};
 use properties::ComputedValues;
 use servo_arc::Arc;
 use shared_lock::{Locked, StylesheetGuards, SharedRwLockReadGuard};
-use stylesheets::{MallocSizeOfFn, PerOrigin, StylesheetContents, StylesheetInDocument};
+use stylesheets::{MallocEnclosingSizeOfFn, MallocSizeOfFn, PerOrigin, StylesheetContents};
+use stylesheets::StylesheetInDocument;
 use stylist::{ExtraStyleData, Stylist};
 
 /// Little wrapper to a Gecko style sheet.
@@ -187,10 +188,17 @@ impl PerDocumentStyleDataImpl {
 
     /// Measures heap usage.
     pub fn malloc_add_size_of_children(&self, malloc_size_of: MallocSizeOfFn,
+                                       malloc_enclosing_size_of: MallocEnclosingSizeOfFn,
                                        sizes: &mut ServoStyleSetSizes) {
-        self.stylist.malloc_add_size_of_children(malloc_size_of, sizes);
+        self.stylist.malloc_add_size_of_children(malloc_size_of, malloc_enclosing_size_of, sizes);
 
-        // We may measure more fields in the future if DMD says it's worth it.
+        let data = &self.extra_style_data;
+        sizes.mStylistOther +=
+            data.user_agent.malloc_size_of_children(malloc_size_of, malloc_enclosing_size_of);
+        sizes.mStylistOther +=
+            data.user.malloc_size_of_children(malloc_size_of, malloc_enclosing_size_of);
+        sizes.mStylistOther +=
+            data.author.malloc_size_of_children(malloc_size_of, malloc_enclosing_size_of);
     }
 }
 

components/style/gecko/generated/bindings.rs

@@ -2063,6 +2063,8 @@ extern "C" {
 }
 extern "C" {
     pub fn Servo_StyleSet_AddSizeOfExcludingThis(malloc_size_of: MallocSizeOf,
+                                                 malloc_enclosing_size_of:
+                                                     MallocSizeOf,
                                                  sizes:
                                                      *mut ServoStyleSetSizes,
                                                  set:

components/style/gecko/generated/structs_debug.rs

@@ -5565,11 +5565,16 @@ pub mod root {
         #[derive(Debug, Copy)]
         pub struct ServoStyleSetSizes {
             pub mStylistRuleTree: usize,
+            pub mStylistPrecomputedPseudos: usize,
+            pub mStylistElementAndPseudosMaps: usize,
+            pub mStylistInvalidationMap: usize,
+            pub mStylistRevalidationSelectors: usize,
+            pub mStylistOther: usize,
             pub mOther: usize,
         }
         #[test]
         fn bindgen_test_layout_ServoStyleSetSizes() {
-            assert_eq!(::std::mem::size_of::<ServoStyleSetSizes>() , 16usize ,
+            assert_eq!(::std::mem::size_of::<ServoStyleSetSizes>() , 56usize ,
                        concat ! (
                        "Size of: " , stringify ! ( ServoStyleSetSizes ) ));
             assert_eq! (::std::mem::align_of::<ServoStyleSetSizes>() , 8usize
@@ -5583,9 +5588,44 @@ pub mod root {
                         "Alignment of field: " , stringify ! (
                         ServoStyleSetSizes ) , "::" , stringify ! (
                         mStylistRuleTree ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistPrecomputedPseudos as * const _ as usize } ,
+                        8usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistPrecomputedPseudos ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistElementAndPseudosMaps as * const _ as usize }
+                        , 16usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistElementAndPseudosMaps ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistInvalidationMap as * const _ as usize } ,
+                        24usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistInvalidationMap ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistRevalidationSelectors as * const _ as usize }
+                        , 32usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistRevalidationSelectors ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistOther as * const _ as usize } , 40usize ,
+                        concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistOther ) ));
             assert_eq! (unsafe {
                         & ( * ( 0 as * const ServoStyleSetSizes ) ) . mOther
-                        as * const _ as usize } , 8usize , concat ! (
+                        as * const _ as usize } , 48usize , concat ! (
                         "Alignment of field: " , stringify ! (
                         ServoStyleSetSizes ) , "::" , stringify ! ( mOther )
                         ));

components/style/gecko/generated/structs_release.rs

@@ -5453,11 +5453,16 @@ pub mod root {
         #[derive(Debug, Copy)]
         pub struct ServoStyleSetSizes {
             pub mStylistRuleTree: usize,
+            pub mStylistPrecomputedPseudos: usize,
+            pub mStylistElementAndPseudosMaps: usize,
+            pub mStylistInvalidationMap: usize,
+            pub mStylistRevalidationSelectors: usize,
+            pub mStylistOther: usize,
             pub mOther: usize,
         }
         #[test]
         fn bindgen_test_layout_ServoStyleSetSizes() {
-            assert_eq!(::std::mem::size_of::<ServoStyleSetSizes>() , 16usize ,
+            assert_eq!(::std::mem::size_of::<ServoStyleSetSizes>() , 56usize ,
                        concat ! (
                        "Size of: " , stringify ! ( ServoStyleSetSizes ) ));
             assert_eq! (::std::mem::align_of::<ServoStyleSetSizes>() , 8usize
@@ -5471,9 +5476,44 @@ pub mod root {
                         "Alignment of field: " , stringify ! (
                         ServoStyleSetSizes ) , "::" , stringify ! (
                         mStylistRuleTree ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistPrecomputedPseudos as * const _ as usize } ,
+                        8usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistPrecomputedPseudos ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistElementAndPseudosMaps as * const _ as usize }
+                        , 16usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistElementAndPseudosMaps ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistInvalidationMap as * const _ as usize } ,
+                        24usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistInvalidationMap ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistRevalidationSelectors as * const _ as usize }
+                        , 32usize , concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistRevalidationSelectors ) ));
+            assert_eq! (unsafe {
+                        & ( * ( 0 as * const ServoStyleSetSizes ) ) .
+                        mStylistOther as * const _ as usize } , 40usize ,
+                        concat ! (
+                        "Alignment of field: " , stringify ! (
+                        ServoStyleSetSizes ) , "::" , stringify ! (
+                        mStylistOther ) ));
             assert_eq! (unsafe {
                         & ( * ( 0 as * const ServoStyleSetSizes ) ) . mOther
-                        as * const _ as usize } , 8usize , concat ! (
+                        as * const _ as usize } , 48usize , concat ! (
                         "Alignment of field: " , stringify ! (
                         ServoStyleSetSizes ) , "::" , stringify ! ( mOther )
                         ));

components/style/invalidation/element/invalidation_map.rs

@@ -14,6 +14,8 @@ use selectors::parser::{Combinator, Component};
 use selectors::parser::{Selector, SelectorIter, SelectorMethods};
 use selectors::visitor::SelectorVisitor;
 use smallvec::SmallVec;
+#[cfg(feature = "gecko")]
+use stylesheets::{MallocEnclosingSizeOfFn, MallocSizeOfHash};
 
 #[cfg(feature = "gecko")]
 /// Gets the element state relevant to the given `:dir` pseudo-class selector.
@@ -137,10 +139,10 @@ impl SelectorMapEntry for StateDependency {
 pub struct InvalidationMap {
     /// A map from a given class name to all the selectors with that class
     /// selector.
-    pub class_to_selector: MaybeCaseInsensitiveHashMap<Atom, SelectorMap<Dependency>>,
+    pub class_to_selector: MaybeCaseInsensitiveHashMap<Atom, SmallVec<[Dependency; 1]>>,
     /// A map from a given id to all the selectors with that ID in the
     /// stylesheets currently applying to the document.
-    pub id_to_selector: MaybeCaseInsensitiveHashMap<Atom, SelectorMap<Dependency>>,
+    pub id_to_selector: MaybeCaseInsensitiveHashMap<Atom, SmallVec<[Dependency; 1]>>,
     /// A map of all the state dependencies.
     pub state_affecting_selectors: SelectorMap<StateDependency>,
     /// A map of other attribute affecting selectors.
@@ -243,21 +245,21 @@ impl InvalidationMap {
             for class in compound_visitor.classes {
                 self.class_to_selector
                     .entry(class, quirks_mode)
-                    .or_insert_with(SelectorMap::new)
+                    .or_insert_with(SmallVec::new)
-                    .insert(Dependency {
+                    .push(Dependency {
                         selector: selector.clone(),
                         selector_offset: sequence_start,
-                    }, quirks_mode);
+                    })
             }
 
             for id in compound_visitor.ids {
                 self.id_to_selector
                     .entry(id, quirks_mode)
-                    .or_insert_with(SelectorMap::new)
+                    .or_insert_with(SmallVec::new)
-                    .insert(Dependency {
+                    .push(Dependency {
                         selector: selector.clone(),
                         selector_offset: sequence_start,
-                    }, quirks_mode);
+                    })
             }
 
             if !compound_visitor.state.is_empty() {
@@ -287,6 +289,24 @@ impl InvalidationMap {
             index += 1; // Account for the combinator.
         }
     }
+
+    /// Measures heap usage.
+    #[cfg(feature = "gecko")]
+    pub fn malloc_size_of_children(&self, malloc_enclosing_size_of: MallocEnclosingSizeOfFn)
+                                   -> usize {
+        // Currently we measure the HashMap storage, but not things pointed to
+        // by keys and values.
+        let mut n = 0;
+
+        n += self.class_to_selector.malloc_shallow_size_of_hash(malloc_enclosing_size_of);
+        n += self.id_to_selector.malloc_shallow_size_of_hash(malloc_enclosing_size_of);
+
+        n += self.state_affecting_selectors.malloc_size_of_children(malloc_enclosing_size_of);
+
+        n += self.other_attribute_affecting_selectors.malloc_size_of_children(
+            malloc_enclosing_size_of);
+        n
+    }
 }
 
 /// A struct that collects invalidations for a given compound selector.

components/style/invalidation/element/invalidator.rs

@@ -23,6 +23,12 @@ use selectors::parser::{Combinator, Component, Selector};
 use smallvec::SmallVec;
 use std::fmt;
 
+#[derive(Debug, PartialEq)]
+enum VisitedDependent {
+    Yes,
+    No,
+}
+
 /// The struct that takes care of encapsulating all the logic on where and how
 /// element styles need to be invalidated.
 pub struct TreeStyleInvalidator<'a, 'b: 'a, E>
@@ -792,20 +798,26 @@ impl<'a, 'b: 'a, E> InvalidationCollector<'a, 'b, E>
         let removed_id = self.removed_id;
         if let Some(ref id) = removed_id {
             if let Some(deps) = map.id_to_selector.get(id, quirks_mode) {
-                self.collect_dependencies_in_map(deps)
+                for dep in deps {
+                    self.scan_dependency(dep, VisitedDependent::No);
+                }
             }
         }
 
         let added_id = self.added_id;
         if let Some(ref id) = added_id {
             if let Some(deps) = map.id_to_selector.get(id, quirks_mode) {
-                self.collect_dependencies_in_map(deps)
+                for dep in deps {
+                    self.scan_dependency(dep, VisitedDependent::No);
+                }
             }
         }
 
         for class in self.classes_added.iter().chain(self.classes_removed.iter()) {
             if let Some(deps) = map.class_to_selector.get(class, quirks_mode) {
-                self.collect_dependencies_in_map(deps)
+                for dep in deps {
+                    self.scan_dependency(dep, VisitedDependent::No);
+                }
             }
         }
 
@@ -839,10 +851,7 @@ impl<'a, 'b: 'a, E> InvalidationCollector<'a, 'b, E>
             self.removed_id,
             self.classes_removed,
             &mut |dependency| {
-                self.scan_dependency(
+                self.scan_dependency(dependency, VisitedDependent::No);
-                    dependency,
-                    /* is_visited_dependent = */ false
-                );
                 true
             },
         );
@@ -862,10 +871,13 @@ impl<'a, 'b: 'a, E> InvalidationCollector<'a, 'b, E>
                 if !dependency.state.intersects(state_changes) {
                     return true;
                 }
-                self.scan_dependency(
+                let visited_dependent =
-                    &dependency.dep,
+                    if dependency.state.intersects(IN_VISITED_OR_UNVISITED_STATE) {
-                    dependency.state.intersects(IN_VISITED_OR_UNVISITED_STATE)
+                        VisitedDependent::Yes
-                );
+                    } else {
+                        VisitedDependent::No
+                    };
+                self.scan_dependency(&dependency.dep, visited_dependent);
                 true
             },
         );
@@ -874,9 +886,9 @@ impl<'a, 'b: 'a, E> InvalidationCollector<'a, 'b, E>
     fn scan_dependency(
         &mut self,
         dependency: &Dependency,
-        is_visited_dependent: bool
+        is_visited_dependent: VisitedDependent,
     ) {
-        debug!("TreeStyleInvalidator::scan_dependency({:?}, {:?}, {})",
+        debug!("TreeStyleInvalidator::scan_dependency({:?}, {:?}, {:?})",
                self.element,
                dependency,
                is_visited_dependent);
@@ -943,7 +955,7 @@ impl<'a, 'b: 'a, E> InvalidationCollector<'a, 'b, E>
         //
         // NOTE: This thing is actually untested because testing it is flaky,
         // see the tests that were added and then backed out in bug 1328509.
-        if is_visited_dependent && now_context.relevant_link_found {
+        if is_visited_dependent == VisitedDependent::Yes && now_context.relevant_link_found {
             then_context.visited_handling = VisitedHandlingMode::RelevantLinkVisited;
             let matched_then =
                 matches_selector(&dependency.selector,

components/style/invalidation/media_queries.rs

@@ -8,8 +8,8 @@ use context::QuirksMode;
 use fnv::FnvHashSet;
 use media_queries::Device;
 use shared_lock::SharedRwLockReadGuard;
-use stylesheets::{DocumentRule, ImportRule, MediaRule,  SupportsRule};
+use stylesheets::{DocumentRule, ImportRule, MallocEnclosingSizeOfFn, MallocSizeOfHash, MediaRule};
-use stylesheets::{NestedRuleIterationCondition, Stylesheet};
+use stylesheets::{NestedRuleIterationCondition, Stylesheet, SupportsRule};
 
 /// A key for a given media query result.
 ///
@@ -88,6 +88,12 @@ impl EffectiveMediaQueryResults {
         // because of stylesheet reusing... shrug.
         self.set.insert(item.to_media_list_key());
     }
+
+    /// Measure heap usage.
+    pub fn malloc_size_of_children(&self, malloc_enclosing_size_of: MallocEnclosingSizeOfFn)
+                                   -> usize {
+        self.set.malloc_shallow_size_of_hash(malloc_enclosing_size_of)
+    }
 }
 
 /// A filter that filters over effective rules, but allowing all potentially

components/style/rule_tree/mod.rs

@@ -796,7 +796,7 @@ impl RuleNode {
     }
 
     fn malloc_size_of_including_self(&self, malloc_size_of: MallocSizeOfFn) -> usize {
-        let mut n = unsafe { malloc_size_of(self as *const _ as *const _) };
+        let mut n = unsafe { (malloc_size_of.0)(self as *const _ as *const _) };
         for child in self.iter_children() {
             n += unsafe { (*child.ptr()).malloc_size_of_including_self(malloc_size_of) };
         }

components/style/selector_map.rs

@@ -19,6 +19,8 @@ use selectors::matching::{matches_selector, MatchingContext, ElementSelectorFlag
 use selectors::parser::{Component, Combinator, SelectorIter};
 use smallvec::{SmallVec, VecLike};
 use std::hash::{BuildHasherDefault, Hash, Hasher};
+#[cfg(feature = "gecko")]
+use stylesheets::{MallocEnclosingSizeOfFn, MallocSizeOfHash};
 use stylist::Rule;
 
 /// A hasher implementation that doesn't hash anything, because it expects its
@@ -144,6 +146,22 @@ impl<T: 'static> SelectorMap<T> {
     pub fn len(&self) -> usize {
         self.count
     }
+
+    /// Measures heap usage.
+    #[cfg(feature = "gecko")]
+    pub fn malloc_size_of_children(&self, malloc_enclosing_size_of: MallocEnclosingSizeOfFn)
+                                   -> usize {
+        // Currently we measure the HashMap storage, but not things pointed to
+        // by keys and values.
+        let mut n = 0;
+        n += self.id_hash.malloc_shallow_size_of_hash(malloc_enclosing_size_of);
+        n += self.class_hash.malloc_shallow_size_of_hash(malloc_enclosing_size_of);
+        n += self.local_name_hash.malloc_shallow_size_of_hash(malloc_enclosing_size_of);
+
+        // We may measure other fields in the future if DMD says it's worth it.
+
+        n
+    }
 }
 
 impl SelectorMap<Rule> {
@@ -502,3 +520,14 @@ impl<V: 'static> MaybeCaseInsensitiveHashMap<Atom, V> {
         }
     }
 }
+
+#[cfg(feature = "gecko")]
+impl<K, V> MallocSizeOfHash for MaybeCaseInsensitiveHashMap<K, V>
+    where K: PrecomputedHash + Eq + Hash
+{
+    fn malloc_shallow_size_of_hash(&self, malloc_enclosing_size_of: MallocEnclosingSizeOfFn)
+                                   -> usize {
+        self.0.malloc_shallow_size_of_hash(malloc_enclosing_size_of)
+    }
+}
+

components/style/selector_parser.rs

@@ -188,4 +188,9 @@ impl<T> PerPseudoElementMap<T> {
         }
         Ok(self.entries[index].as_mut().unwrap())
     }
+
+    /// Get an iterator for the entries.
+    pub fn iter(&self) -> ::std::slice::Iter<Option<T>> {
+        self.entries.iter()
+    }
 }

components/style/stylesheets/keyframes_rule.rs

@@ -19,7 +19,7 @@ use shared_lock::{DeepCloneParams, DeepCloneWithLock, SharedRwLock, SharedRwLock
 use std::fmt;
 use style_traits::{PARSING_MODE_DEFAULT, ToCss, ParseError, StyleParseError};
 use style_traits::PropertyDeclarationParseError;
-use stylesheets::{CssRuleType, StylesheetContents};
+use stylesheets::{CssRuleType, MallocSizeOfFn, MallocSizeOfVec, StylesheetContents};
 use stylesheets::rule_parser::{VendorPrefix, get_location_with_offset};
 use values::{KeyframesName, serialize_percentage};
 
@@ -442,6 +442,14 @@ impl KeyframesAnimation {
 
         result
     }
+
+    /// Measure heap usage.
+    pub fn malloc_size_of_children(&self, malloc_size_of: MallocSizeOfFn) -> usize {
+        let mut n = 0;
+        n += self.steps.malloc_shallow_size_of_vec(malloc_size_of);
+        n += self.properties_changed.malloc_shallow_size_of_vec(malloc_size_of);
+        n
+    }
 }
 
 /// Parses a keyframes list, like:

components/style/stylesheets/memory.rs

@@ -9,16 +9,25 @@ use gecko_bindings::bindings::Gecko_HaveSeenPtr;
 #[cfg(feature = "gecko")]
 use gecko_bindings::structs::SeenPtrs;
 #[cfg(feature = "gecko")]
+use hash::HashMap;
+#[cfg(feature = "gecko")]
 use servo_arc::Arc;
 use shared_lock::SharedRwLockReadGuard;
+use std::collections::HashSet;
+use std::hash::{BuildHasher, Hash};
 use std::os::raw::c_void;
 
 /// Like gecko_bindings::structs::MallocSizeOf, but without the Option<>
 /// wrapper.
 ///
-/// Note that functions of this type should not be called via
+/// Note that functions of this type should be called via do_malloc_size_of(),
-/// do_malloc_size_of(), rather than directly.
+/// rather than directly.
-pub type MallocSizeOfFn = unsafe extern "C" fn(ptr: *const c_void) -> usize;
+#[derive(Clone, Copy)]
+pub struct MallocSizeOfFn(pub unsafe extern "C" fn(ptr: *const c_void) -> usize);
+
+/// Like MallocSizeOfFn, but can take an interior pointer.
+#[derive(Clone, Copy)]
+pub struct MallocEnclosingSizeOfFn(pub unsafe extern "C" fn(ptr: *const c_void) -> usize);
 
 /// Servo-side counterpart to mozilla::SizeOfState. The only difference is that
 /// this struct doesn't contain the SeenPtrs table, just a pointer to it.
@@ -30,7 +39,7 @@ pub struct SizeOfState {
     pub seen_ptrs: *mut SeenPtrs,
 }
 
-/// Call malloc_size_of on ptr, first checking that the allocation isn't empty.
+/// Check if an allocation is empty.
 pub unsafe fn is_empty<T>(ptr: *const T) -> bool {
     return ptr as usize <= ::std::mem::align_of::<T>();
 }
@@ -40,10 +49,18 @@ pub unsafe fn do_malloc_size_of<T>(malloc_size_of: MallocSizeOfFn, ptr: *const T
     if is_empty(ptr) {
         0
     } else {
-        malloc_size_of(ptr as *const c_void)
+        (malloc_size_of.0)(ptr as *const c_void)
     }
 }
 
+/// Call malloc_enclosing_size_of on ptr, which must not be empty.
+pub unsafe fn do_malloc_enclosing_size_of<T>(
+    malloc_enclosing_size_of: MallocEnclosingSizeOfFn, ptr: *const T) -> usize
+{
+    assert!(!is_empty(ptr));
+    (malloc_enclosing_size_of.0)(ptr as *const c_void)
+}
+
 /// Trait for measuring the size of heap data structures.
 pub trait MallocSizeOf {
     /// Measure the size of any heap-allocated structures that hang off this
@@ -92,7 +109,7 @@ impl<T: MallocSizeOfWithRepeats> MallocSizeOfWithRepeats for Arc<T> {
         let mut n = 0;
         let heap_ptr = self.heap_ptr();
         if unsafe { !is_empty(heap_ptr) && !Gecko_HaveSeenPtr(state.seen_ptrs, heap_ptr) } {
-            n += unsafe { (state.malloc_size_of)(heap_ptr) };
+            n += unsafe { (state.malloc_size_of.0)(heap_ptr) };
             n += (**self).malloc_size_of_children(state);
         }
         n
@@ -110,3 +127,78 @@ impl<T: MallocSizeOfWithGuard> MallocSizeOfWithGuard for Vec<T> {
             |n, elem| n + elem.malloc_size_of_children(guard, malloc_size_of))
     }
 }
+
+/// Trait for measuring the heap usage of a Box<T>.
+pub trait MallocSizeOfBox {
+    /// Measure shallowly the size of the memory used by the T -- anything
+    /// pointed to by the T must be measured separately.
+    fn malloc_shallow_size_of_box(&self, malloc_size_of: MallocSizeOfFn) -> usize;
+}
+
+impl<T> MallocSizeOfBox for Box<T> {
+    fn malloc_shallow_size_of_box(&self, malloc_size_of: MallocSizeOfFn) -> usize {
+        unsafe { do_malloc_size_of(malloc_size_of, &**self as *const T) }
+    }
+}
+
+/// Trait for measuring the heap usage of a vector.
+pub trait MallocSizeOfVec {
+    /// Measure shallowly the size of the memory used by the Vec's elements --
+    /// anything pointed to by the elements must be measured separately, using
+    /// iteration.
+    fn malloc_shallow_size_of_vec(&self, malloc_size_of: MallocSizeOfFn) -> usize;
+}
+
+impl<T> MallocSizeOfVec for Vec<T> {
+    fn malloc_shallow_size_of_vec(&self, malloc_size_of: MallocSizeOfFn) -> usize {
+        unsafe { do_malloc_size_of(malloc_size_of, self.as_ptr()) }
+    }
+}
+
+/// Trait for measuring the heap usage of a hash table.
+pub trait MallocSizeOfHash {
+    /// Measure shallowly the size of the memory used within a hash table --
+    /// anything pointer to by the keys and values must be measured separately,
+    /// using iteration.
+    fn malloc_shallow_size_of_hash(&self, malloc_enclosing_size_of: MallocEnclosingSizeOfFn)
+                                   -> usize;
+}
+
+impl<T, S> MallocSizeOfHash for HashSet<T, S>
+    where T: Eq + Hash,
+          S: BuildHasher
+{
+    fn malloc_shallow_size_of_hash(&self, malloc_enclosing_size_of: MallocEnclosingSizeOfFn)
+                                   -> usize {
+        // The first value from the iterator gives us an interior pointer.
+        // malloc_enclosing_size_of() then gives us the storage size. This
+        // assumes that the HashSet's contents (values and hashes) are all
+        // stored in a single contiguous heap allocation.
+        let mut n = 0;
+        for v in self.iter() {
+            n += unsafe { do_malloc_enclosing_size_of(malloc_enclosing_size_of, v as *const T) };
+            break;
+        }
+        n
+    }
+}
+
+#[cfg(feature = "gecko")]
+impl<K, V, S> MallocSizeOfHash for HashMap<K, V, S>
+    where K: Eq + Hash,
+          S: BuildHasher
+{
+    fn malloc_shallow_size_of_hash(&self, malloc_enclosing_size_of: MallocEnclosingSizeOfFn)
+                                   -> usize {
+        // The first value from the iterator gives us an interior pointer.
+        // malloc_enclosing_size_of() then gives us the storage size. This
+        // assumes that the HashMap's contents (keys, values, and hashes) are
+        // all stored in a single contiguous heap allocation.
+        let mut n = 0;
+        for v in self.values() {
+            n += unsafe { do_malloc_enclosing_size_of(malloc_enclosing_size_of, v as *const V) };
+            break;
+        }
+        n
+    }
+}

components/style/stylesheets/mod.rs

@@ -40,7 +40,8 @@ pub use self::import_rule::ImportRule;
 pub use self::keyframes_rule::KeyframesRule;
 pub use self::loader::StylesheetLoader;
 pub use self::media_rule::MediaRule;
-pub use self::memory::{MallocSizeOf, MallocSizeOfFn, MallocSizeOfWithGuard};
+pub use self::memory::{MallocEnclosingSizeOfFn, MallocSizeOf, MallocSizeOfBox, MallocSizeOfFn};
+pub use self::memory::{MallocSizeOfHash, MallocSizeOfVec, MallocSizeOfWithGuard};
 #[cfg(feature = "gecko")]
 pub use self::memory::{MallocSizeOfWithRepeats, SizeOfState};
 pub use self::namespace_rule::NamespaceRule;

components/style/stylist.rs

@@ -43,7 +43,9 @@ use stylesheet_set::{OriginValidity, SheetRebuildKind, StylesheetSet, Stylesheet
 use stylesheets::{CounterStyleRule, FontFaceRule, FontFeatureValuesRule};
 use stylesheets::{CssRule, Origin, OriginSet, PerOrigin, PerOriginIter};
 #[cfg(feature = "gecko")]
-use stylesheets::{MallocSizeOf, MallocSizeOfFn};
+use stylesheets::{MallocEnclosingSizeOfFn, MallocSizeOf, MallocSizeOfBox, MallocSizeOfFn};
+#[cfg(feature = "gecko")]
+use stylesheets::{MallocSizeOfHash, MallocSizeOfVec};
 use stylesheets::StyleRule;
 use stylesheets::StylesheetInDocument;
 use stylesheets::UserAgentStylesheets;
@@ -360,6 +362,25 @@ impl DocumentCascadeData {
             }
         }
     }
+
+    /// Measures heap usage.
+    #[cfg(feature = "gecko")]
+    pub fn malloc_add_size_of_children(&self, malloc_size_of: MallocSizeOfFn,
+                                       malloc_enclosing_size_of: MallocEnclosingSizeOfFn,
+                                       sizes: &mut ServoStyleSetSizes) {
+        self.per_origin.user_agent.malloc_add_size_of_children(malloc_size_of,
+                                                               malloc_enclosing_size_of, sizes);
+        self.per_origin.user.malloc_add_size_of_children(malloc_size_of,
+                                                         malloc_enclosing_size_of, sizes);
+        self.per_origin.author.malloc_add_size_of_children(malloc_size_of,
+                                                           malloc_enclosing_size_of, sizes);
+
+        for elem in self.precomputed_pseudo_element_decls.iter() {
+            if let Some(ref elem) = *elem {
+                sizes.mStylistPrecomputedPseudos += elem.malloc_shallow_size_of_vec(malloc_size_of);
+            }
+        }
+    }
 }
 
 /// A wrapper over a StylesheetSet that can be `Sync`, since it's only used and
@@ -1565,9 +1586,13 @@ impl Stylist {
     /// Measures heap usage.
     #[cfg(feature = "gecko")]
     pub fn malloc_add_size_of_children(&self, malloc_size_of: MallocSizeOfFn,
+                                       malloc_enclosing_size_of: MallocEnclosingSizeOfFn,
                                        sizes: &mut ServoStyleSetSizes) {
-        // XXX: need to measure other fields
+        self.cascade_data.malloc_add_size_of_children(malloc_size_of, malloc_enclosing_size_of,
+                                                      sizes);
         sizes.mStylistRuleTree += self.rule_tree.malloc_size_of_children(malloc_size_of);
+
+        // We may measure other fields in the future if DMD says it's worth it.
     }
 }
 
@@ -1620,6 +1645,17 @@ impl ExtraStyleData {
             self.counter_styles.clear();
         }
     }
+
+    /// Measure heap usage.
+    #[cfg(feature = "gecko")]
+    pub fn malloc_size_of_children(&self, malloc_size_of: MallocSizeOfFn,
+                                   malloc_enclosing_size_of: MallocEnclosingSizeOfFn) -> usize {
+        let mut n = 0;
+        n += self.font_faces.malloc_shallow_size_of_vec(malloc_size_of);
+        n += self.font_feature_values.malloc_shallow_size_of_vec(malloc_size_of);
+        n += self.counter_styles.malloc_shallow_size_of_hash(malloc_enclosing_size_of);
+        n
+    }
 }
 
 /// SelectorMapEntry implementation for use in our revalidation selector map.
@@ -1913,6 +1949,38 @@ impl CascadeData {
         self.mapped_ids.clear();
         self.selectors_for_cache_revalidation.clear();
     }
+
+    /// Measures heap usage.
+    #[cfg(feature = "gecko")]
+    pub fn malloc_add_size_of_children(&self, malloc_size_of: MallocSizeOfFn,
+                                       malloc_enclosing_size_of: MallocEnclosingSizeOfFn,
+                                       sizes: &mut ServoStyleSetSizes) {
+        sizes.mStylistElementAndPseudosMaps +=
+            self.element_map.malloc_size_of_children(malloc_enclosing_size_of);
+
+        for elem in self.pseudos_map.iter() {
+            if let Some(ref elem) = *elem {
+                sizes.mStylistElementAndPseudosMaps +=
+                    elem.malloc_shallow_size_of_box(malloc_size_of) +
+                    elem.malloc_size_of_children(malloc_enclosing_size_of)
+            }
+        }
+
+        sizes.mStylistOther +=
+            self.animations.malloc_shallow_size_of_hash(malloc_enclosing_size_of);
+        for val in self.animations.values() {
+            sizes.mStylistOther += val.malloc_size_of_children(malloc_size_of);
+        }
+
+        sizes.mStylistInvalidationMap +=
+            self.invalidation_map.malloc_size_of_children(malloc_enclosing_size_of);
+
+        sizes.mStylistRevalidationSelectors +=
+            self.selectors_for_cache_revalidation.malloc_size_of_children(malloc_enclosing_size_of);
+
+        sizes.mStylistOther +=
+            self.effective_media_query_results.malloc_size_of_children(malloc_enclosing_size_of);
+    }
 }
 
 impl Default for CascadeData {

ports/geckolib/glue.rs

@@ -118,8 +118,9 @@ use style::shared_lock::{SharedRwLockReadGuard, StylesheetGuards, ToCssWithGuard
 use style::string_cache::Atom;
 use style::style_adjuster::StyleAdjuster;
 use style::stylesheets::{CssRule, CssRules, CssRuleType, CssRulesHelpers, DocumentRule};
-use style::stylesheets::{FontFeatureValuesRule, ImportRule, KeyframesRule, MallocSizeOfWithGuard};
+use style::stylesheets::{FontFeatureValuesRule, ImportRule, KeyframesRule, MallocEnclosingSizeOfFn};
-use style::stylesheets::{MediaRule, NamespaceRule, Origin, OriginSet, PageRule, SizeOfState, StyleRule};
+use style::stylesheets::{MallocSizeOfFn, MallocSizeOfWithGuard, MediaRule};
+use style::stylesheets::{NamespaceRule, Origin, OriginSet, PageRule, SizeOfState, StyleRule};
 use style::stylesheets::{StylesheetContents, SupportsRule};
 use style::stylesheets::StylesheetLoader as StyleStylesheetLoader;
 use style::stylesheets::keyframes_rule::{Keyframe, KeyframeSelector, KeyframesStepValue};
@@ -775,7 +776,7 @@ pub extern "C" fn Servo_Element_ClearData(element: RawGeckoElementBorrowed) {
 pub extern "C" fn Servo_Element_SizeOfExcludingThisAndCVs(malloc_size_of: GeckoMallocSizeOf,
                                                           seen_ptrs: *mut SeenPtrs,
                                                           element: RawGeckoElementBorrowed) -> usize {
-    let malloc_size_of = malloc_size_of.unwrap();
+    let malloc_size_of = MallocSizeOfFn(malloc_size_of.unwrap());
     let element = GeckoElement(element);
     let borrow = element.borrow_data();
     if let Some(data) = borrow {
@@ -1083,9 +1084,8 @@ pub extern "C" fn Servo_StyleSheet_SizeOfIncludingThis(
 ) -> usize {
     let global_style_data = &*GLOBAL_STYLE_DATA;
     let guard = global_style_data.shared_lock.read();
-    let malloc_size_of = malloc_size_of.unwrap();
+    let malloc_size_of = MallocSizeOfFn(malloc_size_of.unwrap());
-    StylesheetContents::as_arc(&sheet)
+    StylesheetContents::as_arc(&sheet).malloc_size_of_children(&guard, malloc_size_of)
-        .malloc_size_of_children(&guard, malloc_size_of)
 }
 
 #[no_mangle]
@@ -3610,13 +3610,15 @@ pub extern "C" fn Servo_StyleSet_ResolveForDeclarations(
 #[no_mangle]
 pub extern "C" fn Servo_StyleSet_AddSizeOfExcludingThis(
     malloc_size_of: GeckoMallocSizeOf,
+    malloc_enclosing_size_of: GeckoMallocSizeOf,
     sizes: *mut ServoStyleSetSizes,
     raw_data: RawServoStyleSetBorrowed
 ) {
     let data = PerDocumentStyleData::from_ffi(raw_data).borrow_mut();
-    let malloc_size_of = malloc_size_of.unwrap();
+    let malloc_size_of = MallocSizeOfFn(malloc_size_of.unwrap());
+    let malloc_enclosing_size_of = MallocEnclosingSizeOfFn(malloc_enclosing_size_of.unwrap());
     let sizes = unsafe { sizes.as_mut() }.unwrap();
-    data.malloc_add_size_of_children(malloc_size_of, sizes);
+    data.malloc_add_size_of_children(malloc_size_of, malloc_enclosing_size_of, sizes);
 }
 
 #[no_mangle]