Add supertraits method to rustc_middle

2 files changed, 50 insertions, 0 deletions

diff --git a/compiler/rustc_middle/src/traits/mod.rs b/compiler/rustc_middle/src/traits/mod.rs
index 49071e7995b..49a64cb246a 100644
--- a/compiler/rustc_middle/src/traits/mod.rs
+++ b/compiler/rustc_middle/src/traits/mod.rs
@@ -7,6 +7,7 @@ pub mod query;
 pub mod select;
 pub mod specialization_graph;
 mod structural_impls;
+pub mod util;
 
 use crate::infer::canonical::Canonical;
 use crate::thir::abstract_const::NotConstEvaluatable;
diff --git a/compiler/rustc_middle/src/traits/util.rs b/compiler/rustc_middle/src/traits/util.rs
new file mode 100644
index 00000000000..3490c688170
--- /dev/null
+++ b/compiler/rustc_middle/src/traits/util.rs
@@ -0,0 +1,49 @@
+use rustc_data_structures::stable_set::FxHashSet;
+
+use crate::ty::{PolyTraitRef, TyCtxt};
+
+/// Given a PolyTraitRef, get the PolyTraitRefs of the trait's (transitive) supertraits.
+///
+/// A simplfied version of the same function at `rustc_infer::traits::util::supertraits`.
+pub fn supertraits<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    trait_ref: PolyTraitRef<'tcx>,
+) -> impl Iterator<Item = PolyTraitRef<'tcx>> {
+    Elaborator { tcx, visited: FxHashSet::from_iter([trait_ref]), stack: vec![trait_ref] }
+}
+
+struct Elaborator<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    visited: FxHashSet<PolyTraitRef<'tcx>>,
+    stack: Vec<PolyTraitRef<'tcx>>,
+}
+
+impl<'tcx> Elaborator<'tcx> {
+    fn elaborate(&mut self, trait_ref: PolyTraitRef<'tcx>) {
+        let supertrait_refs = self
+            .tcx
+            .super_predicates_of(trait_ref.def_id())
+            .predicates
+            .into_iter()
+            .flat_map(|(pred, _)| {
+                pred.subst_supertrait(self.tcx, &trait_ref).to_opt_poly_trait_ref()
+            })
+            .map(|t| t.value)
+            .filter(|supertrait_ref| self.visited.insert(*supertrait_ref));
+
+        self.stack.extend(supertrait_refs);
+    }
+}
+
+impl<'tcx> Iterator for Elaborator<'tcx> {
+    type Item = PolyTraitRef<'tcx>;
+
+    fn next(&mut self) -> Option<PolyTraitRef<'tcx>> {
+        if let Some(trait_ref) = self.stack.pop() {
+            self.elaborate(trait_ref);
+            Some(trait_ref)
+        } else {
+            None
+        }
+    }
+}