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 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 @@ fn super_fold(&self, folder: &mut V) -> ControlFlow { }) } } + +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 { + ControlFlow::Continue(match ty.kind() { + TyKind::Param(p) => self[p], + _ => *ty, + }) + } + + fn fold_const(&mut self, c: &Const) -> ControlFlow { + 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 fn body(&self) -> Body { #[derive(Clone, Debug)] pub struct GenericArgs(pub Vec); +impl std::ops::Index for GenericArgs { + type Output = Ty; + + fn index(&self, index: ParamTy) -> &Self::Output { + self.0[index.index as usize].expect_ty() + } +} + +impl std::ops::Index 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 @@ 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:?}"), },