Add discriminant_for_variant to AdtDef

Author: NotLebedev

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();

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)

compiler/rustc_smir/src/stable_mir/ty.rs

@@ -818,6 +818,15 @@ impl AdtDef {
     pub fn variant(&self, idx: VariantIdx) -> Option<VariantDef> {
         (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.