Avoid unnecessary reallocations in domtree::with_function() (#1011)

cranelift/codegen/src/dominator_tree.rs

@@ -226,7 +226,13 @@ impl DominatorTree {
 
     /// Allocate and compute a dominator tree.
     pub fn with_function(func: &Function, cfg: &ControlFlowGraph) -> Self {
-        let mut domtree = Self::new();
+        let ebb_capacity = func.layout.ebb_capacity();
+        let mut domtree = Self {
+            nodes: SecondaryMap::with_capacity(ebb_capacity),
+            postorder: Vec::with_capacity(ebb_capacity),
+            stack: Vec::new(),
+            valid: false,
+        };
         domtree.compute(func, cfg);
         domtree
     }

cranelift/codegen/src/ir/layout.rs

@@ -60,6 +60,11 @@ impl Layout {
         self.first_ebb = None;
         self.last_ebb = None;
     }
+
+    /// Returns the capacity of the `EbbData` map.
+    pub fn ebb_capacity(&self) -> usize {
+        self.ebbs.capacity()
+    }
 }
 
 /// Sequence numbers.

cranelift/entity/src/map.rs

@@ -52,6 +52,20 @@ where
         }
     }
 
+    /// Create a new, empty map with the specified capacity.
+    ///
+    /// The map will be able to hold exactly `capacity` elements without reallocating.
+    pub fn with_capacity(capacity: usize) -> Self
+    where
+        V: Default,
+    {
+        Self {
+            elems: Vec::with_capacity(capacity),
+            default: Default::default(),
+            unused: PhantomData,
+        }
+    }
+
     /// Create a new empty map with a specified default value.
     ///
     /// This constructor does not require V to implement Default.
@@ -63,6 +77,11 @@ where
         }
     }
 
+    /// Returns the number of elements the map can hold without reallocating.
+    pub fn capacity(&self) -> usize {
+        self.elems.capacity()
+    }
+
     /// Get the element at `k` if it exists.
     pub fn get(&self, k: K) -> Option<&V> {
         self.elems.get(k.index())