Refactor rule tree children

We move the data structure to its own module for better
encapsulation of unsafe code.

components/style/rule_tree/map.rs

@@ -0,0 +1,168 @@
+/* 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/. */
+
+#![forbid(unsafe_code)]
+
+use fxhash::FxHashMap;
+use malloc_size_of::{MallocShallowSizeOf, MallocSizeOfOps};
+use std::hash::Hash;
+use std::mem;
+
+pub(super) struct Map<K, V> {
+    inner: MapInner<K, V>,
+}
+
+enum MapInner<K, V> {
+    Empty,
+    One(V),
+    Map(Box<FxHashMap<K, V>>),
+}
+
+pub(super) struct MapIter<'a, K, V> {
+    inner: MapIterInner<'a, K, V>,
+}
+
+enum MapIterInner<'a, K, V> {
+    One(std::option::IntoIter<&'a V>),
+    Map(std::collections::hash_map::Values<'a, K, V>),
+}
+
+impl<K, V> Default for Map<K, V> {
+    fn default() -> Self {
+        Map {
+            inner: MapInner::Empty,
+        }
+    }
+}
+
+impl<'a, K, V> IntoIterator for &'a Map<K, V> {
+    type Item = &'a V;
+    type IntoIter = MapIter<'a, K, V>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        MapIter {
+            inner: match &self.inner {
+                MapInner::Empty => MapIterInner::One(None.into_iter()),
+                MapInner::One(one) => MapIterInner::One(Some(one).into_iter()),
+                MapInner::Map(map) => MapIterInner::Map(map.values()),
+            },
+        }
+    }
+}
+
+impl<'a, K, V> Iterator for MapIter<'a, K, V> {
+    type Item = &'a V;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match &mut self.inner {
+            MapIterInner::One(one_iter) => one_iter.next(),
+            MapIterInner::Map(map_iter) => map_iter.next(),
+        }
+    }
+}
+
+impl<K, V> Map<K, V>
+where
+    K: Eq + Hash,
+{
+    pub(super) fn is_empty(&self) -> bool {
+        match &self.inner {
+            MapInner::Empty => true,
+            MapInner::One(_) => false,
+            MapInner::Map(map) => map.is_empty(),
+        }
+    }
+
+    #[cfg(debug_assertions)]
+    pub(super) fn len(&self) -> usize {
+        match &self.inner {
+            MapInner::Empty => 0,
+            MapInner::One(_) => 1,
+            MapInner::Map(map) => map.len(),
+        }
+    }
+
+    pub(super) fn get(&self, key: &K, key_from_value: impl FnOnce(&V) -> K) -> Option<&V> {
+        match &self.inner {
+            MapInner::One(one) if *key == key_from_value(one) => Some(one),
+            MapInner::Map(map) => map.get(key),
+            MapInner::Empty | MapInner::One(_) => None,
+        }
+    }
+
+    pub(super) fn get_or_insert_with(
+        &mut self,
+        key: K,
+        key_from_value: impl FnOnce(&V) -> K,
+        new_value: impl FnOnce() -> V,
+    ) -> &mut V {
+        match self.inner {
+            MapInner::Empty => {
+                self.inner = MapInner::One(new_value());
+                match &mut self.inner {
+                    MapInner::One(one) => one,
+                    _ => unreachable!(),
+                }
+            },
+            MapInner::One(_) => {
+                let one = match mem::replace(&mut self.inner, MapInner::Empty) {
+                    MapInner::One(one) => one,
+                    _ => unreachable!(),
+                };
+                // If this panics, the child `one` will be lost.
+                let one_key = key_from_value(&one);
+                // Same for the equality test.
+                if key == one_key {
+                    self.inner = MapInner::One(one);
+                    match &mut self.inner {
+                        MapInner::One(one) => return one,
+                        _ => unreachable!(),
+                    }
+                }
+                self.inner = MapInner::Map(Box::new(FxHashMap::with_capacity_and_hasher(
+                    2,
+                    Default::default(),
+                )));
+                let map = match &mut self.inner {
+                    MapInner::Map(map) => map,
+                    _ => 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)
+            },
+            MapInner::Map(ref mut map) => map.entry(key).or_insert_with(new_value),
+        }
+    }
+
+    pub(super) fn remove(&mut self, key: &K, key_from_value: impl FnOnce(&V) -> K) -> Option<V> {
+        match &mut self.inner {
+            MapInner::One(one) if *key == key_from_value(one) => {
+                match mem::replace(&mut self.inner, MapInner::Empty) {
+                    MapInner::One(one) => Some(one),
+                    _ => unreachable!(),
+                }
+            },
+            MapInner::Map(map) => map.remove(key),
+            MapInner::Empty | MapInner::One(_) => None,
+        }
+    }
+}
+
+impl<K, V> MallocShallowSizeOf for Map<K, V>
+where
+    K: Eq + Hash,
+{
+    fn shallow_size_of(&self, ops: &mut MallocSizeOfOps) -> usize {
+        match &self.inner {
+            MapInner::Map(m) => {
+                // We want to account for both the box and the hashmap.
+                m.shallow_size_of(ops) + (**m).shallow_size_of(ops)
+            },
+            MapInner::One(_) | MapInner::Empty => 0,
+        }
+    }
+}