3 files changed, 92 insertions, 4 deletions

diff --git a/compiler/rustc_smir/src/stable_mir/fold.rs b/compiler/rustc_smir/src/stable_mir/fold.rs
index 3567a5c5a3d..831cfb40a15 100644
--- a/compiler/rustc_smir/src/stable_mir/fold.rs
+++ b/compiler/rustc_smir/src/stable_mir/fold.rs
@@ -3,8 +3,8 @@ 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,
+    Allocation, Binder, Const, ConstDef, ConstantKind, ExistentialPredicate, FnSig, GenericArgKind,
+    GenericArgs, Promoted, RigidTy, TermKind, Ty, TyKind, UnevaluatedConst,
 };
 
 pub trait Folder: Sized {
@@ -205,3 +205,26 @@ impl Foldable for FnSig {
         })
     }
 }
+
+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(),
+        })
+    }
+}
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
index ad7a67275ee..7e344dc516b 100644
--- a/compiler/rustc_smir/src/stable_mir/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -138,6 +138,22 @@ pub struct ImplDef(pub(crate) DefId);
 #[derive(Clone, Debug)]
 pub struct GenericArgs(pub Vec<GenericArgKind>);
 
+impl std::ops::Index<ParamTy> for GenericArgs {
+    type Output = Ty;
+
+    fn index(&self, index: ParamTy) -> &Self::Output {
+        self.0[index.index as usize].expect_ty()
+    }
+}
+
+impl std::ops::Index<ParamConst> for GenericArgs {
+    type Output = Const;
+
+    fn index(&self, index: ParamConst) -> &Self::Output {
+        self.0[index.index as usize].expect_const()
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum GenericArgKind {
     Lifetime(Region),
@@ -145,6 +161,28 @@ pub enum GenericArgKind {
     Const(Const),
 }
 
+impl GenericArgKind {
+    /// Panic if this generic argument is not a type, otherwise
+    /// return the type.
+    #[track_caller]
+    pub fn expect_ty(&self) -> &Ty {
+        match self {
+            GenericArgKind::Type(ty) => ty,
+            _ => panic!("{self:?}"),
+        }
+    }
+
+    /// Panic if this generic argument is not a const, otherwise
+    /// return the const.
+    #[track_caller]
+    pub fn expect_const(&self) -> &Const {
+        match self {
+            GenericArgKind::Const(c) => c,
+            _ => panic!("{self:?}"),
+        }
+    }
+}
+
 #[derive(Clone, Debug)]
 pub enum TermKind {
     Type(Ty),
diff --git a/tests/ui-fulldeps/stable-mir/crate-info.rs b/tests/ui-fulldeps/stable-mir/crate-info.rs
index 31ad42b431d..d55eae86f07 100644
--- a/tests/ui-fulldeps/stable-mir/crate-info.rs
+++ b/tests/ui-fulldeps/stable-mir/crate-info.rs
@@ -8,6 +8,7 @@
 
 #![feature(rustc_private)]
 #![feature(assert_matches)]
+#![feature(control_flow_enum)]
 
 extern crate rustc_hir;
 extern crate rustc_middle;
@@ -15,7 +16,10 @@ extern crate rustc_smir;
 
 use rustc_hir::def::DefKind;
 use rustc_middle::ty::TyCtxt;
-use rustc_smir::{rustc_internal, stable_mir};
+use rustc_smir::{
+    rustc_internal,
+    stable_mir::{self, fold::Foldable},
+};
 use std::assert_matches::assert_matches;
 use std::io::Write;
 use std::ops::ControlFlow;
@@ -116,7 +120,30 @@ fn test_stable_mir(tcx: TyCtxt<'_>) -> ControlFlow<()> {
             stable_mir::mir::Terminator::Call { func, .. } => match func {
                 stable_mir::mir::Operand::Constant(c) => match &c.literal.literal {
                     stable_mir::ty::ConstantKind::Allocated(alloc) => {
-                        assert!(alloc.bytes.is_empty())
+                        assert!(alloc.bytes.is_empty());
+                        match c.literal.ty.kind() {
+                            stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::FnDef(
+                                def,
+                                mut args,
+                            )) => {
+                                let func = def.body();
+                                match func.locals[1]
+                                    .fold(&mut args)
+                                    .continue_value()
+                                    .unwrap()
+                                    .kind()
+                                {
+                                    stable_mir::ty::TyKind::RigidTy(
+                                        stable_mir::ty::RigidTy::Uint(_),
+                                    ) => {}
+                                    stable_mir::ty::TyKind::RigidTy(
+                                        stable_mir::ty::RigidTy::Tuple(_),
+                                    ) => {}
+                                    other => panic!("{other:?}"),
+                                }
+                            }
+                            other => panic!("{other:?}"),
+                        }
                     }
                     other => panic!("{other:?}"),
                 },