Add type folder to SMIR

5 files changed, 255 insertions, 6 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 52ba4bd4e57..671928a63c2 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -103,8 +103,13 @@ impl<'tcx> Context for Tables<'tcx> {
     }
 
     fn ty_kind(&mut self, ty: crate::stable_mir::ty::Ty) -> TyKind {
-        let ty = self.types[ty.0];
-        ty.stable(self)
+        self.types[ty.0].clone().stable(self)
+    }
+
+    fn mk_ty(&mut self, kind: TyKind) -> stable_mir::ty::Ty {
+        let n = self.types.len();
+        self.types.push(MaybeStable::Stable(kind));
+        stable_mir::ty::Ty(n)
     }
 
     fn generics_of(&mut self, def_id: stable_mir::DefId) -> stable_mir::ty::Generics {
@@ -128,20 +133,47 @@ impl<'tcx> Context for Tables<'tcx> {
     }
 }
 
+#[derive(Clone)]
+pub enum MaybeStable<S, R> {
+    Stable(S),
+    Rustc(R),
+}
+
+impl<'tcx, S, R> MaybeStable<S, R> {
+    fn stable(self, tables: &mut Tables<'tcx>) -> S
+    where
+        R: Stable<'tcx, T = S>,
+    {
+        match self {
+            MaybeStable::Stable(s) => s,
+            MaybeStable::Rustc(r) => r.stable(tables),
+        }
+    }
+}
+
+impl<S, R: PartialEq> PartialEq<R> for MaybeStable<S, R> {
+    fn eq(&self, other: &R) -> bool {
+        match self {
+            MaybeStable::Stable(_) => false,
+            MaybeStable::Rustc(r) => r == other,
+        }
+    }
+}
+
 pub struct Tables<'tcx> {
     pub tcx: TyCtxt<'tcx>,
     pub def_ids: Vec<DefId>,
     pub alloc_ids: Vec<AllocId>,
-    pub types: Vec<Ty<'tcx>>,
+    pub types: Vec<MaybeStable<stable_mir::ty::TyKind, Ty<'tcx>>>,
 }
 
 impl<'tcx> Tables<'tcx> {
     fn intern_ty(&mut self, ty: Ty<'tcx>) -> stable_mir::ty::Ty {
-        if let Some(id) = self.types.iter().position(|&t| t == ty) {
+        if let Some(id) = self.types.iter().position(|t| *t == ty) {
             return stable_mir::ty::Ty(id);
         }
         let id = self.types.len();
-        self.types.push(ty);
+        self.types.push(MaybeStable::Rustc(ty));
         stable_mir::ty::Ty(id)
     }
 }
diff --git a/compiler/rustc_smir/src/stable_mir/fold.rs b/compiler/rustc_smir/src/stable_mir/fold.rs
new file mode 100644
index 00000000000..560a99cb497
--- /dev/null
+++ b/compiler/rustc_smir/src/stable_mir/fold.rs
@@ -0,0 +1,207 @@
+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 Folder: Sized {
+    type Break;
+    fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<Self::Break, Ty> {
+        ty.super_fold(self)
+    }
+    fn fold_const(&mut self, c: &Const) -> ControlFlow<Self::Break, Const> {
+        c.super_fold(self)
+    }
+}
+
+pub trait Foldable: Sized + Clone {
+    fn fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        self.super_fold(folder)
+    }
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self>;
+}
+
+impl Foldable for Ty {
+    fn fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        folder.visit_ty(self)
+    }
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        let mut kind = self.kind();
+        match &mut kind {
+            super::ty::TyKind::RigidTy(ty) => *ty = ty.fold(folder)?,
+            super::ty::TyKind::Alias(_, alias) => alias.args = alias.args.fold(folder)?,
+            super::ty::TyKind::Param(_) => {}
+            super::ty::TyKind::Bound(_, _) => {}
+        }
+        ControlFlow::Continue(kind.into())
+    }
+}
+
+impl Foldable for Const {
+    fn fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        folder.fold_const(self)
+    }
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        let mut this = self.clone();
+        match &mut this.literal {
+            super::ty::ConstantKind::Allocated(alloc) => *alloc = alloc.fold(folder)?,
+            super::ty::ConstantKind::Unevaluated(uv) => *uv = uv.fold(folder)?,
+            super::ty::ConstantKind::ParamCt(param) => *param = param.fold(folder)?,
+        }
+        ControlFlow::Continue(this)
+    }
+}
+
+impl Foldable for Opaque {
+    fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(self.clone())
+    }
+}
+
+impl Foldable for Allocation {
+    fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(self.clone())
+    }
+}
+
+impl Foldable for UnevaluatedConst {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        let UnevaluatedConst { ty, def, args, promoted } = self;
+        ControlFlow::Continue(UnevaluatedConst {
+            ty: ty.fold(folder)?,
+            def: def.fold(folder)?,
+            args: args.fold(folder)?,
+            promoted: promoted.fold(folder)?,
+        })
+    }
+}
+
+impl Foldable for ConstDef {
+    fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(self.clone())
+    }
+}
+
+impl<T: Foldable> Foldable for Option<T> {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(match self {
+            Some(val) => Some(val.fold(folder)?),
+            None => None,
+        })
+    }
+}
+
+impl Foldable for Promoted {
+    fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(self.clone())
+    }
+}
+
+impl Foldable for GenericArgs {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(GenericArgs(self.0.fold(folder)?))
+    }
+}
+
+impl Foldable for GenericArgKind {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        let mut this = self.clone();
+        match &mut this {
+            GenericArgKind::Lifetime(lt) => *lt = lt.fold(folder)?,
+            GenericArgKind::Type(t) => *t = t.fold(folder)?,
+            GenericArgKind::Const(c) => *c = c.fold(folder)?,
+        }
+        ControlFlow::Continue(this)
+    }
+}
+
+impl Foldable for RigidTy {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        let mut this = self.clone();
+        match &mut this {
+            RigidTy::Bool
+            | RigidTy::Char
+            | RigidTy::Int(_)
+            | RigidTy::Uint(_)
+            | RigidTy::Float(_)
+            | RigidTy::Never
+            | RigidTy::Foreign(_)
+            | RigidTy::Str => {}
+            RigidTy::Array(t, c) => {
+                *t = t.fold(folder)?;
+                *c = c.fold(folder)?;
+            }
+            RigidTy::Slice(inner) => *inner = inner.fold(folder)?,
+            RigidTy::RawPtr(ty, _) => *ty = ty.fold(folder)?,
+            RigidTy::Ref(_, ty, _) => *ty = ty.fold(folder)?,
+            RigidTy::FnDef(_, args) => *args = args.fold(folder)?,
+            RigidTy::FnPtr(sig) => *sig = sig.fold(folder)?,
+            RigidTy::Closure(_, args) => *args = args.fold(folder)?,
+            RigidTy::Generator(_, args, _) => *args = args.fold(folder)?,
+            RigidTy::Dynamic(pred, r, _) => {
+                *pred = pred.fold(folder)?;
+                *r = r.fold(folder)?;
+            }
+            RigidTy::Tuple(fields) => *fields = fields.fold(folder)?,
+            RigidTy::Adt(_, args) => *args = args.fold(folder)?,
+        }
+        ControlFlow::Continue(this)
+    }
+}
+
+impl<T: Foldable> Foldable for Vec<T> {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        let mut this = self.clone();
+        for arg in &mut this {
+            *arg = arg.fold(folder)?;
+        }
+        ControlFlow::Continue(this)
+    }
+}
+
+impl<T: Foldable> Foldable for Binder<T> {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(Self {
+            value: self.value.fold(folder)?,
+            bound_vars: self.bound_vars.clone(),
+        })
+    }
+}
+
+impl Foldable for ExistentialPredicate {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        let mut this = self.clone();
+        match &mut this {
+            ExistentialPredicate::Trait(tr) => tr.generic_args = tr.generic_args.fold(folder)?,
+            ExistentialPredicate::Projection(p) => {
+                p.term = p.term.fold(folder)?;
+                p.generic_args = p.generic_args.fold(folder)?;
+            }
+            ExistentialPredicate::AutoTrait(_) => {}
+        }
+        ControlFlow::Continue(this)
+    }
+}
+
+impl Foldable for TermKind {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(match self {
+            TermKind::Type(t) => TermKind::Type(t.fold(folder)?),
+            TermKind::Const(c) => TermKind::Const(c.fold(folder)?),
+        })
+    }
+}
+
+impl Foldable for FnSig {
+    fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
+        ControlFlow::Continue(Self {
+            inputs_and_output: self.inputs_and_output.fold(folder)?,
+            c_variadic: self.c_variadic,
+            unsafety: self.unsafety,
+            abi: self.abi.clone(),
+        })
+    }
+}
diff --git a/compiler/rustc_smir/src/stable_mir/mir/body.rs b/compiler/rustc_smir/src/stable_mir/mir/body.rs
index 72f719c2a5e..5957ae0ada9 100644
--- a/compiler/rustc_smir/src/stable_mir/mir/body.rs
+++ b/compiler/rustc_smir/src/stable_mir/mir/body.rs
@@ -391,7 +391,7 @@ pub enum Mutability {
     Mut,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Copy, Clone, Debug)]
 pub enum Safety {
     Unsafe,
     Normal,
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index f9eafd9de7a..90fc569f47b 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -20,6 +20,7 @@ use self::ty::{
 };
 use crate::rustc_smir::Tables;
 
+pub mod fold;
 pub mod mir;
 pub mod ty;
 pub mod visitor;
@@ -158,6 +159,9 @@ pub trait Context {
     /// Obtain the representation of a type.
     fn ty_kind(&mut self, ty: Ty) -> TyKind;
 
+    /// Create a new `Ty` from scratch without information from rustc.
+    fn mk_ty(&mut self, kind: TyKind) -> Ty;
+
     /// HACK: Until we have fully stable consumers, we need an escape hatch
     /// to get `DefId`s out of `CrateItem`s.
     fn rustc_tables(&mut self, f: &mut dyn FnMut(&mut Tables<'_>));
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
index 1db6b1e3d28..5289e233a31 100644
--- a/compiler/rustc_smir/src/stable_mir/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -10,6 +10,12 @@ impl Ty {
     }
 }
 
+impl From<TyKind> for Ty {
+    fn from(value: TyKind) -> Self {
+        with(|context| context.mk_ty(value))
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Const {
     pub literal: ConstantKind,