Create a SMIR visitor

3 files changed, 202 insertions, 3 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 5334a75dc06..d8e1f0877bf 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -5,7 +5,6 @@
 
 use std::fmt::Debug;
 use std::ops::Index;
-use std::string::ToString;
 
 use crate::rustc_internal;
 use crate::{
@@ -156,10 +155,23 @@ pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
 }
 
 /// A type that provides internal information but that can still be used for debug purpose.
-pub type Opaque = impl Debug + ToString + Clone;
+#[derive(Clone)]
+pub struct Opaque(String);
+
+impl std::fmt::Display for Opaque {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+impl std::fmt::Debug for Opaque {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:?}", self.0)
+    }
+}
 
 pub(crate) fn opaque<T: Debug>(value: &T) -> Opaque {
-    format!("{value:?}")
+    Opaque(format!("{value:?}"))
 }
 
 pub struct StableMir {
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 1b834628175..2d43e62e97a 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -20,6 +20,7 @@ use crate::rustc_smir::Tables;
 
 pub mod mir;
 pub mod ty;
+pub mod visitor;
 
 /// Use String for now but we should replace it.
 pub type Symbol = String;
diff --git a/compiler/rustc_smir/src/stable_mir/visitor.rs b/compiler/rustc_smir/src/stable_mir/visitor.rs
new file mode 100644
index 00000000000..c928eb1381f
--- /dev/null
+++ b/compiler/rustc_smir/src/stable_mir/visitor.rs
@@ -0,0 +1,186 @@
+use std::ops::ControlFlow;
+
+use crate::rustc_internal::Opaque;
+
+use super::ty::{
+    Allocation, Binder, Const, ConstDef, ExistentialPredicate, FnSig, GenericArgKind, GenericArgs,
+    Promoted, RigidTy, TermKind, Ty, UnevaluatedConst,
+};
+
+pub trait Visitor: Sized {
+    type Break;
+    fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<Self::Break> {
+        ty.super_visit(self)
+    }
+    fn visit_const(&mut self, c: &Const) -> ControlFlow<Self::Break> {
+        c.super_visit(self)
+    }
+}
+
+pub trait Visitable {
+    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        self.super_visit(visitor)
+    }
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break>;
+}
+
+impl Visitable for Ty {
+    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        visitor.visit_ty(self)
+    }
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        match self.kind() {
+            super::ty::TyKind::RigidTy(ty) => ty.visit(visitor),
+            super::ty::TyKind::Alias(_, alias) => alias.args.visit(visitor),
+            super::ty::TyKind::Param(_) => todo!(),
+            super::ty::TyKind::Bound(_, _) => todo!(),
+        }
+    }
+}
+
+impl Visitable for Const {
+    fn visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        visitor.visit_const(self)
+    }
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        match &self.literal {
+            super::ty::ConstantKind::Allocated(alloc) => alloc.visit(visitor),
+            super::ty::ConstantKind::Unevaluated(uv) => uv.visit(visitor),
+            super::ty::ConstantKind::ParamCt(param) => param.visit(visitor),
+        }
+    }
+}
+
+impl Visitable for Opaque {
+    fn super_visit<V: Visitor>(&self, _visitor: &mut V) -> ControlFlow<V::Break> {
+        ControlFlow::Continue(())
+    }
+}
+
+impl Visitable for Allocation {
+    fn super_visit<V: Visitor>(&self, _visitor: &mut V) -> ControlFlow<V::Break> {
+        ControlFlow::Continue(())
+    }
+}
+
+impl Visitable for UnevaluatedConst {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        let UnevaluatedConst { ty, def, args, promoted } = self;
+        ty.visit(visitor)?;
+        def.visit(visitor)?;
+        args.visit(visitor)?;
+        promoted.visit(visitor)
+    }
+}
+
+impl Visitable for ConstDef {
+    fn super_visit<V: Visitor>(&self, _visitor: &mut V) -> ControlFlow<V::Break> {
+        ControlFlow::Continue(())
+    }
+}
+
+impl<T: Visitable> Visitable for Option<T> {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        match self {
+            Some(val) => val.visit(visitor),
+            None => ControlFlow::Continue(()),
+        }
+    }
+}
+
+impl Visitable for Promoted {
+    fn super_visit<V: Visitor>(&self, _visitor: &mut V) -> ControlFlow<V::Break> {
+        ControlFlow::Continue(())
+    }
+}
+
+impl Visitable for GenericArgs {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        self.0.visit(visitor)
+    }
+}
+
+impl Visitable for GenericArgKind {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        match self {
+            GenericArgKind::Lifetime(lt) => lt.visit(visitor),
+            GenericArgKind::Type(t) => t.visit(visitor),
+            GenericArgKind::Const(c) => c.visit(visitor),
+        }
+    }
+}
+
+impl Visitable for RigidTy {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        match self {
+            RigidTy::Bool
+            | RigidTy::Char
+            | RigidTy::Int(_)
+            | RigidTy::Uint(_)
+            | RigidTy::Float(_)
+            | RigidTy::Never
+            | RigidTy::Foreign(_)
+            | RigidTy::Str => ControlFlow::Continue(()),
+            RigidTy::Array(t, c) => {
+                t.visit(visitor)?;
+                c.visit(visitor)
+            }
+            RigidTy::Slice(inner) => inner.visit(visitor),
+            RigidTy::RawPtr(ty, _) => ty.visit(visitor),
+            RigidTy::Ref(_, ty, _) => ty.visit(visitor),
+            RigidTy::FnDef(_, args) => args.visit(visitor),
+            RigidTy::FnPtr(sig) => sig.visit(visitor),
+            RigidTy::Closure(_, args) => args.visit(visitor),
+            RigidTy::Generator(_, args, _) => args.visit(visitor),
+            RigidTy::Dynamic(pred, r, _) => {
+                pred.visit(visitor)?;
+                r.visit(visitor)
+            }
+            RigidTy::Tuple(fields) => fields.visit(visitor),
+            RigidTy::Adt(_, args) => args.visit(visitor),
+        }
+    }
+}
+
+impl<T: Visitable> Visitable for Vec<T> {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        for arg in self {
+            arg.visit(visitor)?;
+        }
+        ControlFlow::Continue(())
+    }
+}
+
+impl<T: Visitable> Visitable for Binder<T> {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        self.value.visit(visitor)
+    }
+}
+
+impl Visitable for ExistentialPredicate {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        match self {
+            ExistentialPredicate::Trait(tr) => tr.generic_args.visit(visitor),
+            ExistentialPredicate::Projection(p) => {
+                p.term.visit(visitor)?;
+                p.generic_args.visit(visitor)
+            }
+            ExistentialPredicate::AutoTrait(_) => ControlFlow::Continue(()),
+        }
+    }
+}
+
+impl Visitable for TermKind {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        match self {
+            TermKind::Type(t) => t.visit(visitor),
+            TermKind::Const(c) => c.visit(visitor),
+        }
+    }
+}
+
+impl Visitable for FnSig {
+    fn super_visit<V: Visitor>(&self, visitor: &mut V) -> ControlFlow<V::Break> {
+        self.inputs_and_output.visit(visitor)
+    }
+}