Add discriminant_for_variant to CoroutineDef

Author: NotLebedev

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

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)

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;