Rollup merge of rust-lang#115532 - oli-obk:smir_folder, r=compiler-errors

Implement SMIR generic parameter instantiation

Also demonstrates the use of it with a test.

This required a few smaller changes that may conflict with `@ericmarkmartin` work, but should be easy to resolve any conflicts on my end if their stuff lands first.

Author: Dylan-DPC

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -77,8 +77,8 @@ impl<'tcx> Context for Tables<'tcx> {
         impl_trait.stable(self)
     }
 
-    fn mir_body(&mut self, item: &stable_mir::CrateItem) -> stable_mir::mir::Body {
-        let def_id = self[item.0];
+    fn mir_body(&mut self, item: stable_mir::DefId) -> stable_mir::mir::Body {
+        let def_id = self[item];
         let mir = self.tcx.optimized_mir(def_id);
         stable_mir::mir::Body {
             blocks: mir
@@ -102,8 +102,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 {
@@ -127,20 +132,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)
     }
 }
@@ -1096,25 +1128,33 @@ impl<'tcx> Stable<'tcx> for ty::Const<'tcx> {
                         tables,
                     ))
                 }
-                ty::ParamCt(param) => stable_mir::ty::ConstantKind::ParamCt(opaque(&param)),
+                ty::ParamCt(param) => stable_mir::ty::ConstantKind::Param(param.stable(tables)),
                 ty::ErrorCt(_) => unreachable!(),
                 ty::InferCt(_) => unreachable!(),
                 ty::BoundCt(_, _) => unimplemented!(),
                 ty::PlaceholderCt(_) => unimplemented!(),
                 ty::Unevaluated(uv) => {
                     stable_mir::ty::ConstantKind::Unevaluated(stable_mir::ty::UnevaluatedConst {
-                        ty: tables.intern_ty(self.ty()),
                         def: tables.const_def(uv.def),
                         args: uv.args.stable(tables),
                         promoted: None,
                     })
                 }
                 ty::ExprCt(_) => unimplemented!(),
             },
+            ty: tables.intern_ty(self.ty()),
         }
     }
 }
 
+impl<'tcx> Stable<'tcx> for ty::ParamConst {
+    type T = stable_mir::ty::ParamConst;
+    fn stable(&self, _: &mut Tables<'tcx>) -> Self::T {
+        use stable_mir::ty::ParamConst;
+        ParamConst { index: self.index, name: self.name.to_string() }
+    }
+}
+
 impl<'tcx> Stable<'tcx> for ty::ParamTy {
     type T = stable_mir::ty::ParamTy;
     fn stable(&self, _: &mut Tables<'tcx>) -> Self::T {
@@ -1183,22 +1223,27 @@ impl<'tcx> Stable<'tcx> for ty::TraitDef {
 }
 
 impl<'tcx> Stable<'tcx> for rustc_middle::mir::ConstantKind<'tcx> {
-    type T = stable_mir::ty::ConstantKind;
+    type T = stable_mir::ty::Const;
 
     fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
         match *self {
-            ConstantKind::Ty(c) => c.stable(tables).literal,
-            ConstantKind::Unevaluated(unev_const, ty) => {
-                stable_mir::ty::ConstantKind::Unevaluated(stable_mir::ty::UnevaluatedConst {
-                    ty: tables.intern_ty(ty),
-                    def: tables.const_def(unev_const.def),
-                    args: unev_const.args.stable(tables),
-                    promoted: unev_const.promoted.map(|u| u.as_u32()),
-                })
-            }
-            ConstantKind::Val(val, ty) => {
-                stable_mir::ty::ConstantKind::Allocated(alloc::new_allocation(ty, val, tables))
-            }
+            ConstantKind::Ty(c) => c.stable(tables),
+            ConstantKind::Unevaluated(unev_const, ty) => stable_mir::ty::Const {
+                literal: stable_mir::ty::ConstantKind::Unevaluated(
+                    stable_mir::ty::UnevaluatedConst {
+                        def: tables.const_def(unev_const.def),
+                        args: unev_const.args.stable(tables),
+                        promoted: unev_const.promoted.map(|u| u.as_u32()),
+                    },
+                ),
+                ty: tables.intern_ty(ty),
+            },
+            ConstantKind::Val(val, ty) => stable_mir::ty::Const {
+                literal: stable_mir::ty::ConstantKind::Allocated(alloc::new_allocation(
+                    ty, val, tables,
+                )),
+                ty: tables.intern_ty(ty),
+            },
         }
     }
 }

compiler/rustc_smir/src/stable_mir/fold.rs

@@ -0,0 +1,230 @@
+use std::ops::ControlFlow;
+
+use crate::rustc_internal::Opaque;
+
+use super::ty::{
+    Allocation, Binder, Const, ConstDef, ConstantKind, ExistentialPredicate, FnSig, GenericArgKind,
+    GenericArgs, Promoted, RigidTy, TermKind, Ty, TyKind, 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::Param(_) => {}
+        }
+        this.ty = this.ty.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 { def, args, promoted } = self;
+        ControlFlow::Continue(UnevaluatedConst {
+            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(),
+        })
+    }
+}
+
+pub enum Never {}
+
+/// In order to instantiate a `Foldable`'s generic parameters with specific arguments,
+/// `GenericArgs` can be used as a `Folder` that replaces all mentions of generic params
+/// with the entries in its list.
+impl Folder for GenericArgs {
+    type Break = Never;
+
+    fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<Self::Break, Ty> {
+        ControlFlow::Continue(match ty.kind() {
+            TyKind::Param(p) => self[p],
+            _ => *ty,
+        })
+    }
+
+    fn fold_const(&mut self, c: &Const) -> ControlFlow<Self::Break, Const> {
+        ControlFlow::Continue(match &c.literal {
+            ConstantKind::Param(p) => self[p.clone()].clone(),
+            _ => c.clone(),
+        })
+    }
+}