Auto merge of rust-lang#137655 - nnethercote:split-edges-iterator, r=nnethercote

bors · bors · commit 385970f0c1fd · 2025-03-09T15:30:47.000Z
Split the `Edges` iterator.

Some nice performance wins here, mostly on the `wg-grammar` benchmark.

r? `@lcnr`
diff --git 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
@@ -23,8 +20,8 @@ pub(crate) type ReverseConstraintGraph = ConstraintGraph<Reverse>;
 /// Marker trait that controls whether a `R1: R2` constraint
 /// represents an edge `R1 -> R2` or `R2 -> R1`.
 pub(crate) trait ConstraintGraphDirection: Copy + 'static {
-    fn start_region(c: &OutlivesConstraint<'_>) -> RegionVid;
-    fn end_region(c: &OutlivesConstraint<'_>) -> RegionVid;
+    fn start_region(sup: RegionVid, sub: RegionVid) -> RegionVid;
+    fn end_region(sup: RegionVid, sub: RegionVid) -> RegionVid;
     fn is_normal() -> bool;
 }
 
@@ -36,12 +33,12 @@ pub(crate) trait ConstraintGraphDirection: Copy + 'static {
 pub(crate) struct Normal;
 
 impl ConstraintGraphDirection for Normal {
-    fn start_region(c: &OutlivesConstraint<'_>) -> RegionVid {
-        c.sup
+    fn start_region(sup: RegionVid, _sub: RegionVid) -> RegionVid {
+        sup
     }
 
-    fn end_region(c: &OutlivesConstraint<'_>) -> RegionVid {
-        c.sub
+    fn end_region(_sup: RegionVid, sub: RegionVid) -> RegionVid {
+        sub
     }
 
     fn is_normal() -> bool {
@@ -57,12 +54,12 @@ impl ConstraintGraphDirection for Normal {
 pub(crate) struct Reverse;
 
 impl ConstraintGraphDirection for Reverse {
-    fn start_region(c: &OutlivesConstraint<'_>) -> RegionVid {
-        c.sub
+    fn start_region(_sup: RegionVid, sub: RegionVid) -> RegionVid {
+        sub
     }
 
-    fn end_region(c: &OutlivesConstraint<'_>) -> RegionVid {
-        c.sup
+    fn end_region(sup: RegionVid, _sub: RegionVid) -> RegionVid {
+        sup
     }
 
     fn is_normal() -> bool {
@@ -84,7 +81,7 @@ impl<D: ConstraintGraphDirection> ConstraintGraph<D> {
         let mut next_constraints = IndexVec::from_elem(None, &set.outlives);
 
         for (idx, constraint) in set.outlives.iter_enumerated().rev() {
-            let head = &mut first_constraints[D::start_region(constraint)];
+            let head = &mut first_constraints[D::start_region(constraint.sup, constraint.sub)];
             let next = &mut next_constraints[idx];
             debug_assert!(next.is_none());
             *next = *head;
@@ -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))
+        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
@@ -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;
@@ -311,9 +311,11 @@ enum RegionRelationCheckResult {
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
-enum Trace<'tcx> {
+enum Trace<'a, 'tcx> {
     StartRegion,
-    FromOutlivesConstraint(OutlivesConstraint<'tcx>),
+    FromGraph(&'a OutlivesConstraint<'tcx>),
+    FromStatic(RegionVid),
+    FromMember(RegionVid, RegionVid, Span),
     NotVisited,
 }
 
@@ -1764,6 +1766,8 @@ impl<'tcx> RegionInferenceContext<'tcx> {
         let mut context = IndexVec::from_elem(Trace::NotVisited, &self.definitions);
         context[from_region] = Trace::StartRegion;
 
+        let fr_static = self.universal_regions().fr_static;
+
         // Use a deque so that we do a breadth-first search. We will
         // stop at the first match, which ought to be the shortest
         // path (fewest constraints).
@@ -1783,13 +1787,39 @@ impl<'tcx> RegionInferenceContext<'tcx> {
             if target_test(r) {
                 let mut result = vec![];
                 let mut p = r;
+                // This loop is cold and runs at the end, which is why we delay
+                // `OutlivesConstraint` construction until now.
                 loop {
-                    match context[p].clone() {
-                        Trace::NotVisited => {
-                            bug!("found unvisited region {:?} on path to {:?}", p, r)
+                    match context[p] {
+                        Trace::FromGraph(c) => {
+                            p = c.sup;
+                            result.push(*c);
                         }
 
-                        Trace::FromOutlivesConstraint(c) => {
+                        Trace::FromStatic(sub) => {
+                            let c = OutlivesConstraint {
+                                sup: fr_static,
+                                sub,
+                                locations: Locations::All(DUMMY_SP),
+                                span: DUMMY_SP,
+                                category: ConstraintCategory::Internal,
+                                variance_info: ty::VarianceDiagInfo::default(),
+                                from_closure: false,
+                            };
+                            p = c.sup;
+                            result.push(c);
+                        }
+
+                        Trace::FromMember(sup, sub, span) => {
+                            let c = OutlivesConstraint {
+                                sup,
+                                sub,
+                                locations: Locations::All(span),
+                                span,
+                                category: ConstraintCategory::OpaqueType,
+                                variance_info: ty::VarianceDiagInfo::default(),
+                                from_closure: false,
+                            };
                             p = c.sup;
                             result.push(c);
                         }
@@ -1798,6 +1828,10 @@ impl<'tcx> RegionInferenceContext<'tcx> {
                             result.reverse();
                             return Some((result, r));
                         }
+
+                        Trace::NotVisited => {
+                            bug!("found unvisited region {:?} on path to {:?}", p, r)
+                        }
                     }
                 }
             }
@@ -1808,45 +1842,42 @@ 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>| {
-                debug_assert_eq!(constraint.sup, r);
-                let sub_region = constraint.sub;
-                if let Trace::NotVisited = context[sub_region] {
-                    context[sub_region] = Trace::FromOutlivesConstraint(constraint);
-                    deque.push_back(sub_region);
+            let mut handle_trace = #[inline(always)]
+            |sub, trace| {
+                if let Trace::NotVisited = context[sub] {
+                    context[sub] = trace;
+                    deque.push_back(sub);
                 }
             };
 
-            // 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 sub in self.constraint_graph.outgoing_edges_from_static() {
+                    handle_trace(sub, Trace::FromStatic(sub));
+                }
+            } 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;
+                    }
+                    debug_assert_eq!(constraint.sup, r);
+                    handle_trace(constraint.sub, Trace::FromGraph(constraint));
                 }
-                handle_constraint(constraint);
             }
 
             // Member constraints can also give rise to `'r: 'x` edges that
             // were not part of the graph initially, so watch out for those.
             // (But they are extremely rare; this loop is very cold.)
             for constraint in self.applied_member_constraints(self.constraint_sccs.scc(r)) {
+                let sub = constraint.min_choice;
                 let p_c = &self.member_constraints[constraint.member_constraint_index];
-                let constraint = OutlivesConstraint {
-                    sup: r,
-                    sub: constraint.min_choice,
-                    locations: Locations::All(p_c.definition_span),
-                    span: p_c.definition_span,
-                    category: ConstraintCategory::OpaqueType,
-                    variance_info: ty::VarianceDiagInfo::default(),
-                    from_closure: false,
-                };
-                handle_constraint(constraint);
+                handle_trace(sub, Trace::FromMember(r, sub, p_c.definition_span));
             }
         }
 
