Preallocate more things (#21964)

# Objective

- Part of #20115

We should try to minimize allocations where its easy to do so.

## Solution

Pre-allocate things upfront where we know their exact resulting size.

## Testing

It'd probably be good to benchmark to see how much (if any) this helps.

Author: ItsDoot

crates/bevy_ecs/src/schedule/graph/dag.rs

@@ -287,7 +287,9 @@ impl<N: GraphNodeId, S: BuildHasher> DagAnalysis<N, S> {
 
         // build a copy of the graph where the nodes and edges appear in topsorted order
         let mut map = <HashMap<_, _>>::with_capacity_and_hasher(n, Default::default());
-        let mut topsorted = DiGraph::<N>::default();
+        let mut topsorted =
+            DiGraph::<N>::with_capacity(topological_order.len(), graph.edge_count());
+
         // iterate nodes in topological order
         for (i, &node) in topological_order.iter().enumerate() {
             map.insert(node, i);
@@ -302,8 +304,8 @@ impl<N: GraphNodeId, S: BuildHasher> DagAnalysis<N, S> {
         let mut connected = HashSet::default();
         let mut disconnected = Vec::default();
         let mut transitive_edges = Vec::default();
-        let mut transitive_reduction = DiGraph::default();
-        let mut transitive_closure = DiGraph::default();
+        let mut transitive_reduction = DiGraph::with_capacity(topsorted.node_count(), 0);
+        let mut transitive_closure = DiGraph::with_capacity(topsorted.node_count(), 0);
 
         let mut visited = FixedBitSet::with_capacity(n);
 
@@ -605,6 +607,7 @@ impl<K: GraphNodeId, V: GraphNodeId, S: BuildHasher> DagGroups<K, V, S> {
             // Remove the key node from the graph.
             flattening.remove_node(N::from(key));
             // Add all previously collected edges.
+            flattening.reserve_edges(temp.len());
             for (a, b) in temp.drain(..) {
                 flattening.add_edge(a, b);
             }
@@ -636,20 +639,35 @@ impl<K: GraphNodeId, V: GraphNodeId, S: BuildHasher> DagGroups<K, V, S> {
                 }
                 (Err(lhs_key), Ok(rhs)) => {
                     // Edge from a key node to a value node, expand to all values in the key's group
-                    for &lhs in self.get(&lhs_key).into_iter().flatten() {
+                    let Some(lhs_group) = self.get(&lhs_key) else {
+                        continue;
+                    };
+                    flattened.reserve_edges(lhs_group.len());
+                    for &lhs in lhs_group {
                         flattened.add_edge(lhs, rhs);
                     }
                 }
                 (Ok(lhs), Err(rhs_key)) => {
                     // Edge from a value node to a key node, expand to all values in the key's group
-                    for &rhs in self.get(&rhs_key).into_iter().flatten() {
+                    let Some(rhs_group) = self.get(&rhs_key) else {
+                        continue;
+                    };
+                    flattened.reserve_edges(rhs_group.len());
+                    for &rhs in rhs_group {
                         flattened.add_edge(lhs, rhs);
                     }
                 }
                 (Err(lhs_key), Err(rhs_key)) => {
                     // Edge between two key nodes, expand to all combinations of their value nodes
-                    for &lhs in self.get(&lhs_key).into_iter().flatten() {
-                        for &rhs in self.get(&rhs_key).into_iter().flatten() {
+                    let Some(lhs_group) = self.get(&lhs_key) else {
+                        continue;
+                    };
+                    let Some(rhs_group) = self.get(&rhs_key) else {
+                        continue;
+                    };
+                    flattened.reserve_edges(lhs_group.len() * rhs_group.len());
+                    for &lhs in lhs_group {
+                        for &rhs in rhs_group {
                             flattened.add_edge(lhs, rhs);
                         }
                     }

crates/bevy_ecs/src/schedule/graph/graph_map.rs

@@ -208,6 +208,18 @@ impl<const DIRECTED: bool, N: GraphNodeId, S: BuildHasher> Graph<DIRECTED, N, S>
         self.edges.contains(&Self::edge_key(a, b))
     }
 
+    /// Reserve capacity for at least `additional` more nodes to be inserted
+    /// in the graph.
+    pub fn reserve_nodes(&mut self, additional: usize) {
+        self.nodes.reserve(additional);
+    }
+
+    /// Reserve capacity for at least `additional` more edges to be inserted
+    /// in the graph.
+    pub fn reserve_edges(&mut self, additional: usize) {
+        self.edges.reserve(additional);
+    }
+
     /// Return an iterator over the nodes of the graph.
     pub fn nodes(&self) -> impl DoubleEndedIterator<Item = N> + ExactSizeIterator<Item = N> + '_ {
         self.nodes.keys().copied()
@@ -402,7 +414,7 @@ impl<N: GraphNodeId, S: BuildHasher> DiGraph<N, S> {
 
         while let Some(mut scc) = sccs.pop() {
             // only look at nodes and edges in this strongly-connected component
-            let mut subgraph = DiGraph::<N>::default();
+            let mut subgraph = DiGraph::<N>::with_capacity(scc.len(), 0);
             for &node in &scc {
                 subgraph.add_node(node);
             }

crates/bevy_ecs/src/schedule/schedule.rs

@@ -842,6 +842,8 @@ impl ScheduleGraph {
                             &previous_result.nodes
                         };
 
+                        self.dependency
+                            .reserve_edges(previous_nodes.len() * current_nodes.len());
                         for previous_node in previous_nodes {
                             for current_node in current_nodes {
                                 self.dependency.add_edge(*previous_node, *current_node);
@@ -935,6 +937,8 @@ impl ScheduleGraph {
         let in_nodes: Vec<_> = self.hierarchy.neighbors_directed(node, Incoming).collect();
         let out_nodes: Vec<_> = self.hierarchy.neighbors_directed(node, Outgoing).collect();
 
+        self.hierarchy
+            .reserve_edges(in_nodes.len() * out_nodes.len());
         for &in_node in &in_nodes {
             for &out_node in &out_nodes {
                 self.hierarchy.add_edge(in_node, out_node);
@@ -944,6 +948,8 @@ impl ScheduleGraph {
         let in_nodes: Vec<_> = self.dependency.neighbors_directed(node, Incoming).collect();
         let out_nodes: Vec<_> = self.dependency.neighbors_directed(node, Outgoing).collect();
 
+        self.dependency
+            .reserve_edges(in_nodes.len() * out_nodes.len());
         for &in_node in &in_nodes {
             for &out_node in &out_nodes {
                 self.dependency.add_edge(in_node, out_node);