Split the `Edges` iterator.

The `Edges` iterator returns `OutlivesConstraint` elements, which are 72
bytes. This is big enough to affect performance. Return
`&OutlivesConstraint` would be better. However, each `Edges` iterator is
really one of two different iterators. The "from graph" case does a
graph traversal and could return `&OutlivesConstraint`. But the "from
static" case just does a `0..n` iteration and constructs a new
`OutlivesConstraint` from that, so it can't return a reference.

This commit splits `Edges into `EdgesFromGraph` and `EdgesFromStatic`,
which allows them to have different return types. This is a perf win for
the `wg-grammar` benchmark.

2 files changed, 85 insertions, 61 deletions

diff --git a/compiler/rustc_borrowck/src/constraints/graph.rs b/compiler/rustc_borrowck/src/constraints/graph.rs
index 3f2064680ae..1209d8bf596 100644
--- a/compiler/rustc_borrowck/src/constraints/graph.rs
+++ b/compiler/rustc_borrowck/src/constraints/graph.rs
@@ -1,11 +1,8 @@
 use rustc_data_structures::graph;
 use rustc_index::IndexVec;
-use rustc_middle::mir::ConstraintCategory;
-use rustc_middle::ty::{RegionVid, VarianceDiagInfo};
-use rustc_span::DUMMY_SP;
+use rustc_middle::ty::RegionVid;
 
 use crate::constraints::{OutlivesConstraint, OutlivesConstraintIndex, OutlivesConstraintSet};
-use crate::type_check::Locations;
 
 /// The construct graph organizes the constraints by their end-points.
 /// It can be used to view a `R1: R2` constraint as either an edge `R1
@@ -105,63 +102,57 @@ impl<D: ConstraintGraphDirection> ConstraintGraph<D> {
         RegionGraph::new(set, self, static_region)
     }
 
+    pub(crate) fn is_normal(&self) -> bool {
+        D::is_normal()
+    }
+
     /// Given a region `R`, iterate over all constraints `R: R1`.
-    pub(crate) fn outgoing_edges<'a, 'tcx>(
+    pub(crate) fn outgoing_edges_from_graph<'a, 'tcx>(
         &'a self,
         region_sup: RegionVid,
         constraints: &'a OutlivesConstraintSet<'tcx>,
-        static_region: RegionVid,
-    ) -> Edges<'a, 'tcx, D> {
-        //if this is the `'static` region and the graph's direction is normal,
-        //then setup the Edges iterator to return all regions #53178
-        if region_sup == static_region && D::is_normal() {
-            Edges {
-                graph: self,
-                constraints,
-                pointer: None,
-                next_static_idx: Some(0),
-                static_region,
-            }
-        } else {
-            //otherwise, just setup the iterator as normal
-            let first = self.first_constraints[region_sup];
-            Edges { graph: self, constraints, pointer: first, next_static_idx: None, static_region }
-        }
+    ) -> EdgesFromGraph<'a, 'tcx, D> {
+        EdgesFromGraph { graph: self, constraints, pointer: self.first_constraints[region_sup] }
+    }
+
+    /// Returns all regions (#53178).
+    pub(crate) fn outgoing_edges_from_static(&self) -> EdgesFromStatic {
+        EdgesFromStatic { next_static_idx: 0, end_static_idx: self.first_constraints.len() }
     }
 }
 
-pub(crate) struct Edges<'a, 'tcx, D: ConstraintGraphDirection> {
+pub(crate) struct EdgesFromGraph<'a, 'tcx, D: ConstraintGraphDirection> {
     graph: &'a ConstraintGraph<D>,
     constraints: &'a OutlivesConstraintSet<'tcx>,
     pointer: Option<OutlivesConstraintIndex>,
-    next_static_idx: Option<usize>,
-    static_region: RegionVid,
 }
 
-impl<'a, 'tcx, D: ConstraintGraphDirection> Iterator for Edges<'a, 'tcx, D> {
-    type Item = OutlivesConstraint<'tcx>;
+impl<'a, 'tcx, D: ConstraintGraphDirection> Iterator for EdgesFromGraph<'a, 'tcx, D> {
+    type Item = &'a OutlivesConstraint<'tcx>;
 
     fn next(&mut self) -> Option<Self::Item> {
         if let Some(p) = self.pointer {
             self.pointer = self.graph.next_constraints[p];
+            Some(&self.constraints[p])
+        } else {
+            None
+        }
+    }
+}
+
+pub(crate) struct EdgesFromStatic {
+    next_static_idx: usize,
+    end_static_idx: usize,
+}
+
+impl Iterator for EdgesFromStatic {
+    type Item = RegionVid;
 
-            Some(self.constraints[p])
-        } else if let Some(next_static_idx) = self.next_static_idx {
-            self.next_static_idx = if next_static_idx == (self.graph.first_constraints.len() - 1) {
-                None
-            } else {
-                Some(next_static_idx + 1)
-            };
-
-            Some(OutlivesConstraint {
-                sup: self.static_region,
-                sub: next_static_idx.into(),
-                locations: Locations::All(DUMMY_SP),
-                span: DUMMY_SP,
-                category: ConstraintCategory::Internal,
-                variance_info: VarianceDiagInfo::default(),
-                from_closure: false,
-            })
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.next_static_idx < self.end_static_idx {
+            let ret = RegionVid::from_usize(self.next_static_idx);
+            self.next_static_idx += 1;
+            Some(ret)
         } else {
             None
         }
@@ -193,21 +184,38 @@ impl<'a, 'tcx, D: ConstraintGraphDirection> RegionGraph<'a, 'tcx, D> {
     /// Given a region `R`, iterate over all regions `R1` such that
     /// there exists a constraint `R: R1`.
     pub(crate) fn outgoing_regions(&self, region_sup: RegionVid) -> Successors<'a, 'tcx, D> {
-        Successors {
-            edges: self.constraint_graph.outgoing_edges(region_sup, self.set, self.static_region),
+        // If this is the `'static` region and the graph's direction is normal,
+        // then setup the Edges iterator to return all regions (#53178).
+        if region_sup == self.static_region && D::is_normal() {
+            Successors::FromStatic(self.constraint_graph.outgoing_edges_from_static())
+        } else {
+            // Otherwise, just setup the iterator as normal.
+            Successors::FromGraph(
+                self.constraint_graph.outgoing_edges_from_graph(region_sup, self.set),
+            )
         }
     }
 }
 
-pub(crate) struct Successors<'a, 'tcx, D: ConstraintGraphDirection> {
-    edges: Edges<'a, 'tcx, D>,
+pub(crate) enum Successors<'a, 'tcx, D: ConstraintGraphDirection> {
+    FromStatic(EdgesFromStatic),
+    FromGraph(EdgesFromGraph<'a, 'tcx, D>),
 }
 
 impl<'a, 'tcx, D: ConstraintGraphDirection> Iterator for Successors<'a, 'tcx, D> {
     type Item = RegionVid;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.edges.next().map(|c| D::end_region(c.sup, c.sub))
+        match self {
+            Successors::FromStatic(edges) => {
+                // No `D::end_region` call needed here: static successors are only possible when
+                // the direction is `Normal`, so we can directly use what would be the `sub` value.
+                edges.next()
+            }
+            Successors::FromGraph(edges) => {
+                edges.next().map(|constraint| D::end_region(constraint.sup, constraint.sub))
+            }
+        }
     }
 }
 
diff --git a/compiler/rustc_borrowck/src/region_infer/mod.rs b/compiler/rustc_borrowck/src/region_infer/mod.rs
index a00fce08c9b..abe9efe8364 100644
--- a/compiler/rustc_borrowck/src/region_infer/mod.rs
+++ b/compiler/rustc_borrowck/src/region_infer/mod.rs
@@ -21,8 +21,8 @@ use rustc_middle::traits::{ObligationCause, ObligationCauseCode};
 use rustc_middle::ty::fold::fold_regions;
 use rustc_middle::ty::{self, RegionVid, Ty, TyCtxt, TypeFoldable, UniverseIndex};
 use rustc_mir_dataflow::points::DenseLocationMap;
-use rustc_span::Span;
 use rustc_span::hygiene::DesugaringKind;
+use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument, trace};
 
 use crate::BorrowckInferCtxt;
@@ -1809,27 +1809,43 @@ impl<'tcx> RegionInferenceContext<'tcx> {
             // A constraint like `'r: 'x` can come from our constraint
             // graph.
             let fr_static = self.universal_regions().fr_static;
-            let outgoing_edges_from_graph =
-                self.constraint_graph.outgoing_edges(r, &self.constraints, fr_static);
 
             // Always inline this closure because it can be hot.
             let mut handle_constraint = #[inline(always)]
-            |constraint: OutlivesConstraint<'tcx>| {
+            |constraint: &OutlivesConstraint<'tcx>| {
                 debug_assert_eq!(constraint.sup, r);
                 let sub_region = constraint.sub;
                 if let Trace::NotVisited = context[sub_region] {
-                    context[sub_region] = Trace::FromOutlivesConstraint(constraint);
+                    context[sub_region] = Trace::FromOutlivesConstraint(*constraint);
                     deque.push_back(sub_region);
                 }
             };
 
-            // This loop can be hot.
-            for constraint in outgoing_edges_from_graph {
-                if matches!(constraint.category, ConstraintCategory::IllegalUniverse) {
-                    debug!("Ignoring illegal universe constraint: {constraint:?}");
-                    continue;
+            // If this is the `'static` region and the graph's direction is normal, then set up the
+            // Edges iterator to return all regions (#53178).
+            if r == fr_static && self.constraint_graph.is_normal() {
+                for next_static_idx in self.constraint_graph.outgoing_edges_from_static() {
+                    let constraint = OutlivesConstraint {
+                        sup: fr_static,
+                        sub: next_static_idx,
+                        locations: Locations::All(DUMMY_SP),
+                        span: DUMMY_SP,
+                        category: ConstraintCategory::Internal,
+                        variance_info: ty::VarianceDiagInfo::default(),
+                        from_closure: false,
+                    };
+                    handle_constraint(&constraint);
+                }
+            } else {
+                let edges = self.constraint_graph.outgoing_edges_from_graph(r, &self.constraints);
+                // This loop can be hot.
+                for constraint in edges {
+                    if matches!(constraint.category, ConstraintCategory::IllegalUniverse) {
+                        debug!("Ignoring illegal universe constraint: {constraint:?}");
+                        continue;
+                    }
+                    handle_constraint(constraint);
                 }
-                handle_constraint(constraint);
             }
 
             // Member constraints can also give rise to `'r: 'x` edges that
@@ -1846,7 +1862,7 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                     variance_info: ty::VarianceDiagInfo::default(),
                     from_closure: false,
                 };
-                handle_constraint(constraint);
+                handle_constraint(&constraint);
             }
         }