Use associated types for SCC annotations, per code review suggestion

1 files changed, 33 insertions, 37 deletions

diff --git a/compiler/rustc_data_structures/src/graph/scc/mod.rs b/compiler/rustc_data_structures/src/graph/scc/mod.rs
index c0fe9e8ff93..518817ea0f5 100644
--- a/compiler/rustc_data_structures/src/graph/scc/mod.rs
+++ b/compiler/rustc_data_structures/src/graph/scc/mod.rs
@@ -10,6 +10,7 @@
 
 use std::assert_matches::debug_assert_matches;
 use std::fmt::Debug;
+use std::marker::PhantomData;
 use std::ops::Range;
 
 use rustc_index::{Idx, IndexSlice, IndexVec};
@@ -49,16 +50,21 @@ pub trait Annotation: Debug + Copy {
 }
 
 /// An accumulator for annotations.
-pub trait Annotations<N: Idx, S: Idx, A: Annotation> {
-    fn new(&self, element: N) -> A;
-    fn annotate_scc(&mut self, scc: S, annotation: A);
+pub trait Annotations<N: Idx> {
+    type Ann: Annotation;
+    type SccIdx: Idx + Ord;
+
+    fn new(&self, element: N) -> Self::Ann;
+    fn annotate_scc(&mut self, scc: Self::SccIdx, annotation: Self::Ann);
 }
 
 /// The nil annotation accumulator, which does nothing.
-impl<N: Idx, S: Idx> Annotations<N, S, ()> for () {
-    fn new(&self, _element: N) -> () {
-        ()
-    }
+struct NoAnnotations<S: Idx + Ord>(PhantomData<S>);
+
+impl<N: Idx, S: Idx + Ord> Annotations<N> for NoAnnotations<S> {
+    type SccIdx = S;
+    type Ann = ();
+    fn new(&self, _element: N) {}
     fn annotate_scc(&mut self, _scc: S, _annotation: ()) {}
 }
 
@@ -112,13 +118,13 @@ struct SccData<S: Idx> {
 impl<N: Idx, S: Idx + Ord> Sccs<N, S> {
     /// Compute SCCs without annotations.
     pub fn new(graph: &impl Successors<Node = N>) -> Self {
-        Self::new_with_annotation(graph, &mut ())
+        Self::new_with_annotation(graph, &mut NoAnnotations(PhantomData::<S>))
     }
 
     /// Compute SCCs and annotate them with a user-supplied annotation
-    pub fn new_with_annotation<A: Annotation, AA: Annotations<N, S, A>>(
+    pub fn new_with_annotation<A: Annotations<N, SccIdx = S>>(
         graph: &impl Successors<Node = N>,
-        annotations: &mut AA,
+        annotations: &mut A,
     ) -> Self {
         SccsConstruction::construct(graph, annotations)
     }
@@ -141,13 +147,7 @@ impl<N: Idx, S: Idx + Ord> Sccs<N, S> {
     pub fn all_sccs(&self) -> impl Iterator<Item = S> + 'static {
         (0..self.scc_data.len()).map(S::new)
     }
-    /*
-        /// Returns an iterator over the SCC annotations in the graph
-        /// The order is the same as `all_sccs()`, dependency order.
-        pub fn all_annotations(&self, annotations: &A) -> impl Iterator<Item = (S, A)> + use<'_, N, S, A> {
-            self.all_sccs().map(|scc| (scc, self.annotation(scc)))
-        }
-    */
+
     /// Returns the SCC to which a node `r` belongs.
     pub fn scc(&self, r: N) -> S {
         self.scc_indices[r]
@@ -223,19 +223,17 @@ impl<S: Idx> SccData<S> {
     }
 }
 
-struct SccsConstruction<'c, 'a, G, S, A, AA>
+struct SccsConstruction<'c, 'a, G, A>
 where
     G: DirectedGraph + Successors,
-    S: Idx,
-    A: Annotation,
-    AA: Annotations<G::Node, S, A>,
+    A: Annotations<G::Node>,
 {
     graph: &'c G,
 
     /// The state of each node; used during walk to record the stack
     /// and after walk to record what cycle each node ended up being
     /// in.
-    node_states: IndexVec<G::Node, NodeState<G::Node, S, A>>,
+    node_states: IndexVec<G::Node, NodeState<G::Node, A::SccIdx, A::Ann>>,
 
     /// The stack of nodes that we are visiting as part of the DFS.
     node_stack: Vec<G::Node>,
@@ -244,21 +242,21 @@ where
     /// position in this stack, and when we encounter a successor SCC,
     /// we push it on the stack. When we complete an SCC, we can pop
     /// everything off the stack that was found along the way.
-    successors_stack: Vec<S>,
+    successors_stack: Vec<A::SccIdx>,
 
     /// A set used to strip duplicates. As we accumulate successors
     /// into the successors_stack, we sometimes get duplicate entries.
     /// We use this set to remove those -- we also keep its storage
     /// around between successors to amortize memory allocation costs.
-    duplicate_set: FxHashSet<S>,
+    duplicate_set: FxHashSet<A::SccIdx>,
 
-    scc_data: SccData<S>,
+    scc_data: SccData<A::SccIdx>,
 
-    annotations: &'a mut AA,
+    annotations: &'a mut A,
 }
 
 #[derive(Copy, Clone, Debug)]
-enum NodeState<N, S, A> {
+enum NodeState<N, S, A: Annotation> {
     /// This node has not yet been visited as part of the DFS.
     ///
     /// After SCC construction is complete, this state ought to be
@@ -289,7 +287,7 @@ enum NodeState<N, S, A> {
 
 /// The state of walking a given node.
 #[derive(Copy, Clone, Debug)]
-enum WalkReturn<S, A> {
+enum WalkReturn<S, A: Annotation> {
     /// The walk found a cycle, but the entire component is not known to have
     /// been fully walked yet. We only know the minimum depth of  this
     /// component in a minimum spanning tree of the graph. This component
@@ -302,12 +300,10 @@ enum WalkReturn<S, A> {
     Complete { scc_index: S, annotation: A },
 }
 
-impl<'c, 'a, G, S, A, AA> SccsConstruction<'c, 'a, G, S, A, AA>
+impl<'c, 'a, G, A> SccsConstruction<'c, 'a, G, A>
 where
     G: DirectedGraph + Successors,
-    S: Idx,
-    A: Annotation,
-    AA: Annotations<G::Node, S, A>,
+    A: Annotations<G::Node>,
 {
     /// Identifies SCCs in the graph `G` and computes the resulting
     /// DAG. This uses a variant of [Tarjan's
@@ -323,7 +319,7 @@ where
     /// Additionally, we keep track of a current annotation of the SCC.
     ///
     /// [wikipedia]: https://bit.ly/2EZIx84
-    fn construct(graph: &'c G, annotations: &'a mut AA) -> Sccs<G::Node, S> {
+    fn construct(graph: &'c G, annotations: &'a mut A) -> Sccs<G::Node, A::SccIdx> {
         let num_nodes = graph.num_nodes();
 
         let mut this = Self {
@@ -349,7 +345,7 @@ where
         Sccs { scc_indices, scc_data: this.scc_data }
     }
 
-    fn start_walk_from(&mut self, node: G::Node) -> WalkReturn<S, A> {
+    fn start_walk_from(&mut self, node: G::Node) -> WalkReturn<A::SccIdx, A::Ann> {
         self.inspect_node(node).unwrap_or_else(|| self.walk_unvisited_node(node))
     }
 
@@ -365,7 +361,7 @@ where
     /// Otherwise, we are looking at a node that has already been
     /// completely visited. We therefore return `WalkReturn::Complete`
     /// with its associated SCC index.
-    fn inspect_node(&mut self, node: G::Node) -> Option<WalkReturn<S, A>> {
+    fn inspect_node(&mut self, node: G::Node) -> Option<WalkReturn<A::SccIdx, A::Ann>> {
         Some(match self.find_state(node) {
             NodeState::InCycle { scc_index, annotation } => {
                 WalkReturn::Complete { scc_index, annotation }
@@ -388,7 +384,7 @@ where
     /// of `r2` (and updates `r` to reflect current result). This is
     /// basically the "find" part of a standard union-find algorithm
     /// (with path compression).
-    fn find_state(&mut self, mut node: G::Node) -> NodeState<G::Node, S, A> {
+    fn find_state(&mut self, mut node: G::Node) -> NodeState<G::Node, A::SccIdx, A::Ann> {
         // To avoid recursion we temporarily reuse the `parent` of each
         // InCycleWith link to encode a downwards link while compressing
         // the path. After we have found the root or deepest node being
@@ -511,7 +507,7 @@ where
     /// caller decide avoids mutual recursion between the two methods and allows
     /// us to maintain an allocated stack for nodes on the path between calls.
     #[instrument(skip(self, initial), level = "trace")]
-    fn walk_unvisited_node(&mut self, initial: G::Node) -> WalkReturn<S, A> {
+    fn walk_unvisited_node(&mut self, initial: G::Node) -> WalkReturn<A::SccIdx, A::Ann> {
         trace!("Walk unvisited node: {initial:?}");
         struct VisitingNodeFrame<G: DirectedGraph, Successors, A> {
             node: G::Node,