From b2bfa1049fe6ca6ad903e947547e230ee54b1a34 Mon Sep 17 00:00:00 2001 From: NotLebedev Date: Tue, 27 May 2025 11:12:05 +0300 Subject: [PATCH 1/6] Add discriminant_for_variant to AdtDef --- compiler/rustc_smir/src/rustc_smir/context.rs | 16 ++++++++++++++-- .../src/stable_mir/compiler_interface.rs | 9 +++++++-- compiler/rustc_smir/src/stable_mir/ty.rs | 9 +++++++++ 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs index bac5c9066f1f6..a803961276a4c 100644 --- a/compiler/rustc_smir/src/rustc_smir/context.rs +++ b/compiler/rustc_smir/src/rustc_smir/context.rs @@ -22,9 +22,9 @@ use stable_mir::mir::mono::{InstanceDef, StaticDef}; use stable_mir::mir::{BinOp, Body, Place, UnOp}; use stable_mir::target::{MachineInfo, MachineSize}; use stable_mir::ty::{ - AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, FieldDef, FnDef, ForeignDef, + AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Discr, FieldDef, FnDef, ForeignDef, ForeignItemKind, GenericArgs, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span, Ty, - TyConst, TyKind, UintTy, VariantDef, + TyConst, TyKind, UintTy, VariantDef, VariantIdx, }; use stable_mir::{Crate, CrateDef, CrateItem, CrateNum, DefId, Error, Filename, ItemKind, Symbol}; @@ -440,6 +440,18 @@ impl<'tcx> SmirCtxt<'tcx> { def.internal(&mut *tables, tcx).variants().len() } + /// Discriminant for a given variant index of AdtDef + pub fn adt_discr_for_variant(&self, adt: AdtDef, variant: VariantIdx) -> Discr { + let mut tables = self.0.borrow_mut(); + let tcx = tables.tcx; + + let discr = adt + .internal(&mut *tables, tcx) + .discriminant_for_variant(tcx, variant.internal(&mut *tables, tcx)); + + Discr { val: discr.val, ty: discr.ty.stable(&mut *tables) } + } + /// The name of a variant. pub fn variant_name(&self, def: VariantDef) -> Symbol { let mut tables = self.0.borrow_mut(); diff --git a/compiler/rustc_smir/src/stable_mir/compiler_interface.rs b/compiler/rustc_smir/src/stable_mir/compiler_interface.rs index bb35e23a72884..510ee699be0fc 100644 --- a/compiler/rustc_smir/src/stable_mir/compiler_interface.rs +++ b/compiler/rustc_smir/src/stable_mir/compiler_interface.rs @@ -13,10 +13,10 @@ use stable_mir::mir::mono::{Instance, InstanceDef, StaticDef}; use stable_mir::mir::{BinOp, Body, Place, UnOp}; use stable_mir::target::MachineInfo; use stable_mir::ty::{ - AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, FieldDef, FnDef, ForeignDef, + AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Discr, FieldDef, FnDef, ForeignDef, ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates, Generics, ImplDef, ImplTrait, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span, TraitDecl, - TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef, + TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef, VariantIdx, }; use stable_mir::{ AssocItems, Crate, CrateItem, CrateItems, CrateNum, DefId, Error, Filename, ImplTraitDecls, @@ -225,6 +225,11 @@ impl<'tcx> SmirInterface<'tcx> { self.cx.adt_variants_len(def) } + /// Discriminant for a given variant index of AdtDef + pub(crate) fn adt_discr_for_variant(&self, adt: AdtDef, variant: VariantIdx) -> Discr { + self.cx.adt_discr_for_variant(adt, variant) + } + /// The name of a variant. pub(crate) fn variant_name(&self, def: VariantDef) -> Symbol { self.cx.variant_name(def) diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs index e331e5934716a..d439142f5fde1 100644 --- a/compiler/rustc_smir/src/stable_mir/ty.rs +++ b/compiler/rustc_smir/src/stable_mir/ty.rs @@ -818,6 +818,15 @@ impl AdtDef { pub fn variant(&self, idx: VariantIdx) -> Option { (idx.to_index() < self.num_variants()).then_some(VariantDef { idx, adt_def: *self }) } + + pub fn discriminant_for_variant(&self, idx: VariantIdx) -> Discr { + with(|cx| cx.adt_discr_for_variant(*self, idx)) + } +} + +pub struct Discr { + pub val: u128, + pub ty: Ty, } /// Definition of a variant, which can be either a struct / union field or an enum variant. From a2d5942e3693ca534ba40c28c7c655574ff475d3 Mon Sep 17 00:00:00 2001 From: NotLebedev Date: Tue, 27 May 2025 11:18:37 +0300 Subject: [PATCH 2/6] Add discriminant_for_variant to CoroutineDef --- compiler/rustc_smir/src/rustc_smir/context.rs | 28 ++++++++++++++++--- .../src/stable_mir/compiler_interface.rs | 18 +++++++++--- compiler/rustc_smir/src/stable_mir/ty.rs | 6 ++++ 3 files changed, 44 insertions(+), 8 deletions(-) diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs index a803961276a4c..ebf91a3e370d8 100644 --- a/compiler/rustc_smir/src/rustc_smir/context.rs +++ b/compiler/rustc_smir/src/rustc_smir/context.rs @@ -12,7 +12,8 @@ use rustc_middle::ty::layout::{ }; use rustc_middle::ty::print::{with_forced_trimmed_paths, with_no_trimmed_paths}; use rustc_middle::ty::{ - GenericPredicates, Instance, List, ScalarInt, TyCtxt, TypeVisitableExt, ValTree, + CoroutineArgsExt, GenericPredicates, Instance, List, ScalarInt, TyCtxt, TypeVisitableExt, + ValTree, }; use rustc_middle::{mir, ty}; use rustc_span::def_id::LOCAL_CRATE; @@ -22,9 +23,9 @@ use stable_mir::mir::mono::{InstanceDef, StaticDef}; use stable_mir::mir::{BinOp, Body, Place, UnOp}; use stable_mir::target::{MachineInfo, MachineSize}; use stable_mir::ty::{ - AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Discr, FieldDef, FnDef, ForeignDef, - ForeignItemKind, GenericArgs, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span, Ty, - TyConst, TyKind, UintTy, VariantDef, VariantIdx, + AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, CoroutineDef, Discr, FieldDef, FnDef, + ForeignDef, ForeignItemKind, GenericArgs, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, + Span, Ty, TyConst, TyKind, UintTy, VariantDef, VariantIdx, }; use stable_mir::{Crate, CrateDef, CrateItem, CrateNum, DefId, Error, Filename, ItemKind, Symbol}; @@ -452,6 +453,25 @@ impl<'tcx> SmirCtxt<'tcx> { Discr { val: discr.val, ty: discr.ty.stable(&mut *tables) } } + /// Discriminant for a given variand index and args of a coroutine + pub fn coroutine_discr_for_variant( + &self, + coroutine: CoroutineDef, + args: &GenericArgs, + variant: VariantIdx, + ) -> Discr { + let mut tables = self.0.borrow_mut(); + let tcx = tables.tcx; + + let discr = args.internal(&mut *tables, tcx).as_coroutine().discriminant_for_variant( + coroutine.def_id().internal(&mut *tables, tcx), + tcx, + variant.internal(&mut *tables, tcx), + ); + + Discr { val: discr.val, ty: discr.ty.stable(&mut *tables) } + } + /// The name of a variant. pub fn variant_name(&self, def: VariantDef) -> Symbol { let mut tables = self.0.borrow_mut(); diff --git a/compiler/rustc_smir/src/stable_mir/compiler_interface.rs b/compiler/rustc_smir/src/stable_mir/compiler_interface.rs index 510ee699be0fc..18313f11c460c 100644 --- a/compiler/rustc_smir/src/stable_mir/compiler_interface.rs +++ b/compiler/rustc_smir/src/stable_mir/compiler_interface.rs @@ -13,10 +13,10 @@ use stable_mir::mir::mono::{Instance, InstanceDef, StaticDef}; use stable_mir::mir::{BinOp, Body, Place, UnOp}; use stable_mir::target::MachineInfo; use stable_mir::ty::{ - AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Discr, FieldDef, FnDef, ForeignDef, - ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates, Generics, - ImplDef, ImplTrait, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span, TraitDecl, - TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef, VariantIdx, + AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, CoroutineDef, Discr, FieldDef, FnDef, + ForeignDef, ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates, + Generics, ImplDef, ImplTrait, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span, + TraitDecl, TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef, VariantIdx, }; use stable_mir::{ AssocItems, Crate, CrateItem, CrateItems, CrateNum, DefId, Error, Filename, ImplTraitDecls, @@ -230,6 +230,16 @@ impl<'tcx> SmirInterface<'tcx> { self.cx.adt_discr_for_variant(adt, variant) } + /// Discriminant for a given variand index and args of a coroutine + pub(crate) fn coroutine_discr_for_variant( + &self, + coroutine: CoroutineDef, + args: &GenericArgs, + variant: VariantIdx, + ) -> Discr { + self.cx.coroutine_discr_for_variant(coroutine, args, variant) + } + /// The name of a variant. pub(crate) fn variant_name(&self, def: VariantDef) -> Symbol { self.cx.variant_name(def) diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs index d439142f5fde1..6976a6e8c6bad 100644 --- a/compiler/rustc_smir/src/stable_mir/ty.rs +++ b/compiler/rustc_smir/src/stable_mir/ty.rs @@ -747,6 +747,12 @@ crate_def! { pub CoroutineDef; } +impl CoroutineDef { + pub fn discriminant_for_variant(&self, args: &GenericArgs, idx: VariantIdx) -> Discr { + with(|cx| cx.coroutine_discr_for_variant(*self, args, idx)) + } +} + crate_def! { #[derive(Serialize)] pub CoroutineClosureDef; From 7d13f8d470c25f68e3ba81a930dfccea5a9210bd Mon Sep 17 00:00:00 2001 From: NotLebedev Date: Thu, 27 Mar 2025 22:20:42 +0300 Subject: [PATCH 3/6] Implement Stable for Discr --- compiler/rustc_smir/src/rustc_smir/context.rs | 21 +++++++------------ .../rustc_smir/src/rustc_smir/convert/ty.rs | 8 +++++++ 2 files changed, 15 insertions(+), 14 deletions(-) diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs index ebf91a3e370d8..7d9679cd67d74 100644 --- a/compiler/rustc_smir/src/rustc_smir/context.rs +++ b/compiler/rustc_smir/src/rustc_smir/context.rs @@ -445,12 +445,9 @@ impl<'tcx> SmirCtxt<'tcx> { pub fn adt_discr_for_variant(&self, adt: AdtDef, variant: VariantIdx) -> Discr { let mut tables = self.0.borrow_mut(); let tcx = tables.tcx; - - let discr = adt - .internal(&mut *tables, tcx) - .discriminant_for_variant(tcx, variant.internal(&mut *tables, tcx)); - - Discr { val: discr.val, ty: discr.ty.stable(&mut *tables) } + let adt = adt.internal(&mut *tables, tcx); + let variant = variant.internal(&mut *tables, tcx); + adt.discriminant_for_variant(tcx, variant).stable(&mut *tables) } /// Discriminant for a given variand index and args of a coroutine @@ -462,14 +459,10 @@ impl<'tcx> SmirCtxt<'tcx> { ) -> Discr { let mut tables = self.0.borrow_mut(); let tcx = tables.tcx; - - let discr = args.internal(&mut *tables, tcx).as_coroutine().discriminant_for_variant( - coroutine.def_id().internal(&mut *tables, tcx), - tcx, - variant.internal(&mut *tables, tcx), - ); - - Discr { val: discr.val, ty: discr.ty.stable(&mut *tables) } + let coroutine = coroutine.def_id().internal(&mut *tables, tcx); + let args = args.internal(&mut *tables, tcx); + let variant = variant.internal(&mut *tables, tcx); + args.as_coroutine().discriminant_for_variant(coroutine, tcx, variant).stable(&mut *tables) } /// The name of a variant. diff --git a/compiler/rustc_smir/src/rustc_smir/convert/ty.rs b/compiler/rustc_smir/src/rustc_smir/convert/ty.rs index 8bcac4c4678e4..a5a90549157c0 100644 --- a/compiler/rustc_smir/src/rustc_smir/convert/ty.rs +++ b/compiler/rustc_smir/src/rustc_smir/convert/ty.rs @@ -959,3 +959,11 @@ impl<'tcx> Stable<'tcx> for ty::ImplTraitInTraitData { } } } + +impl<'tcx> Stable<'tcx> for rustc_middle::ty::util::Discr<'tcx> { + type T = stable_mir::ty::Discr; + + fn stable(&self, tables: &mut Tables<'_>) -> Self::T { + stable_mir::ty::Discr { val: self.val, ty: self.ty.stable(tables) } + } +} From becdd21b77398496970951a41492ac236ecdfc9b Mon Sep 17 00:00:00 2001 From: NotLebedev Date: Wed, 4 Jun 2025 13:28:48 +0300 Subject: [PATCH 4/6] Add test for `AdtDef::discriminant_for_variant` --- tests/ui-fulldeps/stable-mir/check_variant.rs | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 tests/ui-fulldeps/stable-mir/check_variant.rs diff --git a/tests/ui-fulldeps/stable-mir/check_variant.rs b/tests/ui-fulldeps/stable-mir/check_variant.rs new file mode 100644 index 0000000000000..61dc3c4e4b396 --- /dev/null +++ b/tests/ui-fulldeps/stable-mir/check_variant.rs @@ -0,0 +1,154 @@ +//@ run-pass +//! Test that users are able to use stable mir APIs to retrieve type information from a crate item +//! definition. + +//@ ignore-stage1 +//@ ignore-cross-compile +//@ ignore-remote +//@ edition: 2024 + +#![feature(rustc_private)] +#![feature(assert_matches)] + +extern crate rustc_middle; +#[macro_use] +extern crate rustc_smir; +extern crate rustc_driver; +extern crate rustc_interface; +extern crate stable_mir; + +use std::io::Write; +use std::ops::ControlFlow; + +use stable_mir::CrateItem; +use stable_mir::crate_def::CrateDef; +use stable_mir::mir::{AggregateKind, Rvalue, Statement, StatementKind}; +use stable_mir::ty::{IntTy, RigidTy, Ty}; + +const CRATE_NAME: &str = "crate_variant_ty"; + +/// Test if we can retrieve discriminant info for different types. +fn test_def_tys() -> ControlFlow<()> { + check_adt_mono(); + check_adt_poly(); + + ControlFlow::Continue(()) +} + +fn check_adt_mono() { + let mono = get_fn("mono").expect_body(); + + check_statement_is_aggregate_assign( + &mono.blocks[0].statements[0], + 0, + RigidTy::Int(IntTy::Isize), + ); + check_statement_is_aggregate_assign( + &mono.blocks[1].statements[0], + 1, + RigidTy::Int(IntTy::Isize), + ); + check_statement_is_aggregate_assign( + &mono.blocks[2].statements[0], + 2, + RigidTy::Int(IntTy::Isize), + ); +} + +fn check_adt_poly() { + let poly = get_fn("poly").expect_body(); + + check_statement_is_aggregate_assign( + &poly.blocks[0].statements[0], + 0, + RigidTy::Int(IntTy::Isize), + ); + check_statement_is_aggregate_assign( + &poly.blocks[1].statements[0], + 1, + RigidTy::Int(IntTy::Isize), + ); + check_statement_is_aggregate_assign( + &poly.blocks[2].statements[0], + 2, + RigidTy::Int(IntTy::Isize), + ); +} + +fn get_fn(name: &str) -> CrateItem { + stable_mir::all_local_items().into_iter().find(|it| it.name().eq(name)).unwrap() +} + +fn check_statement_is_aggregate_assign( + statement: &Statement, + expected_discr_val: u128, + expected_discr_ty: RigidTy, +) { + if let Statement { kind: StatementKind::Assign(_, rvalue), .. } = statement + && let Rvalue::Aggregate(aggregate, _) = rvalue + && let AggregateKind::Adt(adt_def, variant_idx, ..) = aggregate + { + let discr = adt_def.discriminant_for_variant(*variant_idx); + + assert_eq!(discr.val, expected_discr_val); + assert_eq!(discr.ty, Ty::from_rigid_kind(expected_discr_ty)); + } else { + unreachable!("Unexpected statement"); + } +} + +/// This test will generate and analyze a dummy crate using the stable mir. +/// For that, it will first write the dummy crate into a file. +/// Then it will create a `StableMir` using custom arguments and then +/// it will run the compiler. +fn main() { + let path = "defs_ty_input.rs"; + generate_input(&path).unwrap(); + let args = &[ + "rustc".to_string(), + "-Cpanic=abort".to_string(), + "--crate-name".to_string(), + CRATE_NAME.to_string(), + path.to_string(), + ]; + run!(args, test_def_tys).unwrap(); +} + +fn generate_input(path: &str) -> std::io::Result<()> { + let mut file = std::fs::File::create(path)?; + write!( + file, + r#" + use std::hint::black_box; + + enum Mono {{ + A, + B(i32), + C {{ a: i32, b: u32 }}, + }} + + enum Poly {{ + A, + B(T), + C {{ t: T }}, + }} + + pub fn main() {{ + mono(); + }} + + fn mono() {{ + black_box(Mono::A); + black_box(Mono::B(6)); + black_box(Mono::C {{a: 1, b: 10 }}); + }} + + fn poly() {{ + black_box(Poly::::A); + black_box(Poly::B(1i32)); + black_box(Poly::C {{ t: 1i32 }}); + }} + "# + )?; + Ok(()) +} From a39e98fe1da2772baa55b27eef7e56e55ac09675 Mon Sep 17 00:00:00 2001 From: NotLebedev Date: Wed, 4 Jun 2025 14:25:37 +0300 Subject: [PATCH 5/6] Add test for `AdtDef::discriminant_for_variant` polymorphic over parameter --- tests/ui-fulldeps/stable-mir/check_variant.rs | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/ui-fulldeps/stable-mir/check_variant.rs b/tests/ui-fulldeps/stable-mir/check_variant.rs index 61dc3c4e4b396..412fb8be35751 100644 --- a/tests/ui-fulldeps/stable-mir/check_variant.rs +++ b/tests/ui-fulldeps/stable-mir/check_variant.rs @@ -31,6 +31,7 @@ const CRATE_NAME: &str = "crate_variant_ty"; fn test_def_tys() -> ControlFlow<()> { check_adt_mono(); check_adt_poly(); + check_adt_poly2(); ControlFlow::Continue(()) } @@ -75,6 +76,26 @@ fn check_adt_poly() { ); } +fn check_adt_poly2() { + let poly = get_fn("poly2").expect_body(); + + check_statement_is_aggregate_assign( + &poly.blocks[0].statements[0], + 0, + RigidTy::Int(IntTy::Isize), + ); + check_statement_is_aggregate_assign( + &poly.blocks[1].statements[0], + 1, + RigidTy::Int(IntTy::Isize), + ); + check_statement_is_aggregate_assign( + &poly.blocks[2].statements[0], + 2, + RigidTy::Int(IntTy::Isize), + ); +} + fn get_fn(name: &str) -> CrateItem { stable_mir::all_local_items().into_iter().find(|it| it.name().eq(name)).unwrap() } @@ -135,6 +156,8 @@ fn generate_input(path: &str) -> std::io::Result<()> { pub fn main() {{ mono(); + poly(); + poly2::(1); }} fn mono() {{ @@ -148,6 +171,12 @@ fn generate_input(path: &str) -> std::io::Result<()> { black_box(Poly::B(1i32)); black_box(Poly::C {{ t: 1i32 }}); }} + + fn poly2(t: T) {{ + black_box(Poly::::A); + black_box(Poly::B(t)); + black_box(Poly::C {{ t: t }}); + }} "# )?; Ok(()) From 443e45ca01511a4a86314becd443f6735ee2554e Mon Sep 17 00:00:00 2001 From: NotLebedev Date: Fri, 6 Jun 2025 15:25:32 +0300 Subject: [PATCH 6/6] Fix test description --- tests/ui-fulldeps/stable-mir/check_variant.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/ui-fulldeps/stable-mir/check_variant.rs b/tests/ui-fulldeps/stable-mir/check_variant.rs index 412fb8be35751..b0de3369830b2 100644 --- a/tests/ui-fulldeps/stable-mir/check_variant.rs +++ b/tests/ui-fulldeps/stable-mir/check_variant.rs @@ -1,6 +1,6 @@ //@ run-pass -//! Test that users are able to use stable mir APIs to retrieve type information from a crate item -//! definition. +//! Test that users are able to use stable mir APIs to retrieve +//! discriminant value and type for AdtDef and Coroutine variants //@ ignore-stage1 //@ ignore-cross-compile