refactor: rewrite `adt_repr()`, `adt_discr_for_variant()` and `coroutine_discr_for_variant()`

4 files changed, 88 insertions, 19 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/context/impls.rs b/compiler/rustc_smir/src/rustc_smir/context/impls.rs
index edeedbd7089..89ae47143ef 100644
--- a/compiler/rustc_smir/src/rustc_smir/context/impls.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context/impls.rs
@@ -4,13 +4,14 @@
 
 use std::iter;
 
-use rustc_abi::{Endian, Layout};
+use rustc_abi::{Endian, Layout, ReprOptions};
 use rustc_hir::def::DefKind;
 use rustc_hir::{Attribute, LangItem};
 use rustc_middle::mir::interpret::{AllocId, ConstAllocation, ErrorHandled, GlobalAlloc, Scalar};
 use rustc_middle::mir::{BinOp, Body, Const as MirConst, ConstValue, UnOp};
 use rustc_middle::ty::layout::{FnAbiOf, LayoutOf};
 use rustc_middle::ty::print::{with_forced_trimmed_paths, with_no_trimmed_paths};
+use rustc_middle::ty::util::Discr;
 use rustc_middle::ty::{
     AdtDef, AdtKind, AssocItem, Binder, ClosureKind, CoroutineArgsExt, EarlyBinder,
     ExistentialTraitRef, FnSig, GenericArgsRef, Instance, InstanceKind, IntrinsicDef, List,
@@ -363,6 +364,11 @@ impl<'tcx, B: Bridge> SmirCtxt<'tcx, B> {
         self.tcx.is_lang_item(def_id, LangItem::CStr)
     }
 
+    /// Returns the representation options for this ADT.
+    pub fn adt_repr(&self, def: AdtDef<'tcx>) -> ReprOptions {
+        def.repr()
+    }
+
     /// Retrieve the function signature for the given generic arguments.
     pub fn fn_sig(
         &self,
@@ -394,6 +400,25 @@ impl<'tcx, B: Bridge> SmirCtxt<'tcx, B> {
         def.variants().len()
     }
 
+    /// Discriminant for a given variant index of AdtDef.
+    pub fn adt_discr_for_variant(
+        &self,
+        adt: AdtDef<'tcx>,
+        variant: rustc_abi::VariantIdx,
+    ) -> Discr<'tcx> {
+        adt.discriminant_for_variant(self.tcx, variant)
+    }
+
+    /// Discriminant for a given variand index and args of a coroutine.
+    pub fn coroutine_discr_for_variant(
+        &self,
+        coroutine: DefId,
+        args: GenericArgsRef<'tcx>,
+        variant: rustc_abi::VariantIdx,
+    ) -> Discr<'tcx> {
+        args.as_coroutine().discriminant_for_variant(coroutine, self.tcx, variant)
+    }
+
     /// The name of a variant.
     pub fn variant_name(&self, def: &'tcx VariantDef) -> String {
         def.name.to_string()
diff --git a/compiler/rustc_smir/src/stable_mir/compiler_interface.rs b/compiler/rustc_smir/src/stable_mir/compiler_interface.rs
index 3577c15f911..c7f3d25cd71 100644
--- a/compiler/rustc_smir/src/stable_mir/compiler_interface.rs
+++ b/compiler/rustc_smir/src/stable_mir/compiler_interface.rs
@@ -152,6 +152,9 @@ pub(crate) trait SmirInterface {
     /// Returns whether this definition is a C string.
     fn adt_is_cstr(&self, def: AdtDef) -> bool;
 
+    /// Returns the representation options for this ADT.
+    fn adt_repr(&self, def: AdtDef) -> ReprOptions;
+
     /// Retrieve the function signature for the given generic arguments.
     fn fn_sig(&self, def: FnDef, args: &GenericArgs) -> PolyFnSig;
 
@@ -167,6 +170,17 @@ pub(crate) trait SmirInterface {
     /// The number of variants in this ADT.
     fn adt_variants_len(&self, def: AdtDef) -> usize;
 
+    /// Discriminant for a given variant index of AdtDef.
+    fn adt_discr_for_variant(&self, adt: AdtDef, variant: VariantIdx) -> Discr;
+
+    /// Discriminant for a given variand index and args of a coroutine.
+    fn coroutine_discr_for_variant(
+        &self,
+        coroutine: CoroutineDef,
+        args: &GenericArgs,
+        variant: VariantIdx,
+    ) -> Discr;
+
     /// The name of a variant.
     fn variant_name(&self, def: VariantDef) -> Symbol;
     fn variant_fields(&self, def: VariantDef) -> Vec<FieldDef>;
@@ -588,8 +602,10 @@ impl<'tcx> SmirInterface for SmirContainer<'tcx, BridgeTys> {
     }
 
     /// Returns the representation options for this ADT
-    pub(crate) fn adt_repr(&self, def: AdtDef) -> ReprOptions {
-        self.cx.adt_repr(def)
+    fn adt_repr(&self, def: AdtDef) -> ReprOptions {
+        let mut tables = self.tables.borrow_mut();
+        let cx = &*self.cx.borrow();
+        cx.adt_repr(def.internal(&mut *tables, cx.tcx)).stable(&mut *tables, cx)
     }
 
     /// Retrieve the function signature for the given generic arguments.
@@ -632,19 +648,31 @@ impl<'tcx> SmirInterface for SmirContainer<'tcx, BridgeTys> {
         cx.adt_variants_len(def.internal(&mut *tables, cx.tcx))
     }
 
-    /// 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)
+    /// Discriminant for a given variant index of AdtDef.
+    fn adt_discr_for_variant(&self, adt: AdtDef, variant: VariantIdx) -> Discr {
+        let mut tables = self.tables.borrow_mut();
+        let cx = &*self.cx.borrow();
+        cx.adt_discr_for_variant(
+            adt.internal(&mut *tables, cx.tcx),
+            variant.internal(&mut *tables, cx.tcx),
+        )
+        .stable(&mut *tables, cx)
     }
 
-    /// Discriminant for a given variand index and args of a coroutine
-    pub(crate) fn coroutine_discr_for_variant(
+    /// Discriminant for a given variand index and args of a coroutine.
+    fn coroutine_discr_for_variant(
         &self,
         coroutine: CoroutineDef,
         args: &GenericArgs,
         variant: VariantIdx,
     ) -> Discr {
-        self.cx.coroutine_discr_for_variant(coroutine, args, variant)
+        let mut tables = self.tables.borrow_mut();
+        let cx = &*self.cx.borrow();
+        let tcx = cx.tcx;
+        let def = coroutine.def_id().internal(&mut *tables, tcx);
+        let args_ref = args.internal(&mut *tables, tcx);
+        cx.coroutine_discr_for_variant(def, args_ref, variant.internal(&mut *tables, tcx))
+            .stable(&mut *tables, cx)
     }
 
     /// The name of a variant.
diff --git a/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/abi.rs b/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/abi.rs
index 0d8ecf4df2a..d8823a0d10c 100644
--- a/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/abi.rs
+++ b/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/abi.rs
@@ -362,7 +362,11 @@ impl<'tcx> Stable<'tcx> for rustc_abi::WrappingRange {
 impl<'tcx> Stable<'tcx> for rustc_abi::ReprFlags {
     type T = ReprFlags;
 
-    fn stable(&self, _tables: &mut Tables<'_>) -> Self::T {
+    fn stable<'cx>(
+        &self,
+        _tables: &mut Tables<'cx, BridgeTys>,
+        _cx: &SmirCtxt<'cx, BridgeTys>,
+    ) -> Self::T {
         ReprFlags {
             is_simd: self.intersects(Self::IS_SIMD),
             is_c: self.intersects(Self::IS_C),
@@ -375,11 +379,15 @@ impl<'tcx> Stable<'tcx> for rustc_abi::ReprFlags {
 impl<'tcx> Stable<'tcx> for rustc_abi::IntegerType {
     type T = IntegerType;
 
-    fn stable(&self, tables: &mut Tables<'_>) -> Self::T {
+    fn stable<'cx>(
+        &self,
+        tables: &mut Tables<'cx, BridgeTys>,
+        cx: &SmirCtxt<'cx, BridgeTys>,
+    ) -> Self::T {
         match self {
             rustc_abi::IntegerType::Pointer(signed) => IntegerType::Pointer { is_signed: *signed },
             rustc_abi::IntegerType::Fixed(integer, signed) => {
-                IntegerType::Fixed { length: integer.stable(tables), is_signed: *signed }
+                IntegerType::Fixed { length: integer.stable(tables, cx), is_signed: *signed }
             }
         }
     }
@@ -388,12 +396,16 @@ impl<'tcx> Stable<'tcx> for rustc_abi::IntegerType {
 impl<'tcx> Stable<'tcx> for rustc_abi::ReprOptions {
     type T = ReprOptions;
 
-    fn stable(&self, tables: &mut Tables<'_>) -> Self::T {
+    fn stable<'cx>(
+        &self,
+        tables: &mut Tables<'cx, BridgeTys>,
+        cx: &SmirCtxt<'cx, BridgeTys>,
+    ) -> Self::T {
         ReprOptions {
-            int: self.int.map(|int| int.stable(tables)),
-            align: self.align.map(|align| align.stable(tables)),
-            pack: self.pack.map(|pack| pack.stable(tables)),
-            flags: self.flags.stable(tables),
+            int: self.int.map(|int| int.stable(tables, cx)),
+            align: self.align.map(|align| align.stable(tables, cx)),
+            pack: self.pack.map(|pack| pack.stable(tables, cx)),
+            flags: self.flags.stable(tables, cx),
         }
     }
 }
diff --git a/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/ty.rs b/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/ty.rs
index 13f9fd6a86a..c0a430079d8 100644
--- a/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/unstable/convert/stable/ty.rs
@@ -1139,7 +1139,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) }
+    fn stable<'cx>(
+        &self,
+        tables: &mut Tables<'cx, BridgeTys>,
+        cx: &SmirCtxt<'cx, BridgeTys>,
+    ) -> Self::T {
+        stable_mir::ty::Discr { val: self.val, ty: self.ty.stable(tables, cx) }
     }
 }