2 files changed, 66 insertions, 46 deletions

diff --git a/components/style/rule_tree/core.rs b/components/style/rule_tree/core.rs
index ab389aa4e31..e1c6906c053 100644
--- a/components/style/rule_tree/core.rs
+++ b/components/style/rule_tree/core.rs
@@ -13,11 +13,10 @@ use smallvec::SmallVec;
 use std::fmt;
 use std::hash;
 use std::io::Write;
-use std::mem;
 use std::ptr;
 use std::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
 
-use super::map::Map;
+use super::map::{Entry, Map};
 use super::unsafe_box::UnsafeBox;
 use super::{CascadeLevel, StyleSource};
 
@@ -402,33 +401,21 @@ impl StrongRuleNode {
             return child.upgrade();
         }
         let mut children = RwLockUpgradableReadGuard::upgrade(children);
-        let mut is_new = false;
-        let weak = {
-            let is_new = &mut is_new;
-            children.get_or_insert_with(
-                key,
-                |node| node.p.key(),
-                move || {
-                    *is_new = true;
-                    let root = unsafe { root.downgrade() };
-                    let strong = StrongRuleNode::new(Box::new(RuleNode::new(
-                        root,
-                        self.clone(),
-                        source,
-                        level,
-                    )));
-                    let weak = unsafe { strong.downgrade() };
-                    mem::forget(strong);
-                    weak
-                },
-            )
-        };
-
-        if !is_new {
-            return weak.upgrade();
+        match children.entry(key, |node| node.p.key()) {
+            Entry::Occupied(child) => {
+                child.upgrade()
+            },
+            Entry::Vacant(entry) => {
+                let node = StrongRuleNode::new(Box::new(RuleNode::new(
+                    unsafe { root.downgrade() },
+                    self.clone(),
+                    source,
+                    level,
+                )));
+                entry.insert(unsafe { node.downgrade() });
+                node
+            },
         }
-
-        unsafe { StrongRuleNode::from_unsafe_box(UnsafeBox::clone(&weak.p)) }
     }
 
     /// Get the style source corresponding to this rule node. May return `None`
diff --git a/components/style/rule_tree/map.rs b/components/style/rule_tree/map.rs
index 11534d95f7a..33c470e9c12 100644
--- a/components/style/rule_tree/map.rs
+++ b/components/style/rule_tree/map.rs
@@ -6,6 +6,7 @@
 
 use fxhash::FxHashMap;
 use malloc_size_of::{MallocShallowSizeOf, MallocSizeOfOps};
+use std::collections::hash_map;
 use std::hash::Hash;
 use std::mem;
 
@@ -28,6 +29,20 @@ enum MapIterInner<'a, K, V> {
     Map(std::collections::hash_map::Values<'a, K, V>),
 }
 
+pub(super) enum Entry<'a, K, V> {
+    Occupied(&'a mut V),
+    Vacant(VacantEntry<'a, K, V>),
+}
+
+pub(super) struct VacantEntry<'a, K, V> {
+    inner: VacantEntryInner<'a, K, V>,
+}
+
+enum VacantEntryInner<'a, K, V> {
+    One(&'a mut MapInner<K, V>),
+    Map(hash_map::VacantEntry<'a, K, V>),
+}
+
 impl<K, V> Default for Map<K, V> {
     fn default() -> Self {
         Map {
@@ -91,20 +106,15 @@ where
         }
     }
 
-    pub(super) fn get_or_insert_with(
+    pub(super) fn entry(
         &mut self,
         key: K,
         key_from_value: impl FnOnce(&V) -> K,
-        new_value: impl FnOnce() -> V,
-    ) -> &mut V {
+    ) -> Entry<'_, K, V> {
         match self.inner {
-            MapInner::Empty => {
-                self.inner = MapInner::One(new_value());
-                match &mut self.inner {
-                    MapInner::One(one) => one,
-                    _ => unreachable!(),
-                }
-            },
+            ref mut inner @ MapInner::Empty => Entry::Vacant(VacantEntry {
+                inner: VacantEntryInner::One(inner),
+            }),
             MapInner::One(_) => {
                 let one = match mem::replace(&mut self.inner, MapInner::Empty) {
                     MapInner::One(one) => one,
@@ -115,10 +125,11 @@ where
                 // Same for the equality test.
                 if key == one_key {
                     self.inner = MapInner::One(one);
-                    match &mut self.inner {
-                        MapInner::One(one) => return one,
+                    let one = match &mut self.inner {
+                        MapInner::One(one) => one,
                         _ => unreachable!(),
-                    }
+                    };
+                    return Entry::Occupied(one);
                 }
                 self.inner = MapInner::Map(Box::new(FxHashMap::with_capacity_and_hasher(
                     2,
@@ -129,12 +140,19 @@ where
                     _ => unreachable!(),
                 };
                 map.insert(one_key, one);
-                // But it doesn't matter if f panics, by this point
-                // the map is as before but represented as a map instead
-                // of a single value.
-                map.entry(key).or_insert_with(new_value)
+                match map.entry(key) {
+                    hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
+                        inner: VacantEntryInner::Map(entry),
+                    }),
+                    _ => unreachable!(),
+                }
+            },
+            MapInner::Map(ref mut map) => match map.entry(key) {
+                hash_map::Entry::Occupied(entry) => Entry::Occupied(entry.into_mut()),
+                hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry {
+                    inner: VacantEntryInner::Map(entry),
+                }),
             },
-            MapInner::Map(ref mut map) => map.entry(key).or_insert_with(new_value),
         }
     }
 
@@ -152,6 +170,21 @@ where
     }
 }
 
+impl<'a, K, V> VacantEntry<'a, K, V> {
+    pub(super) fn insert(self, value: V) -> &'a mut V {
+        match self.inner {
+            VacantEntryInner::One(map) => {
+                *map = MapInner::One(value);
+                match map {
+                    MapInner::One(one) => one,
+                    _ => unreachable!(),
+                }
+            },
+            VacantEntryInner::Map(entry) => entry.insert(value),
+        }
+    }
+}
+
 impl<K, V> MallocShallowSizeOf for Map<K, V>
 where
     K: Eq + Hash,