Add FieldDef to StableMIR and methods to get type

6 files changed, 95 insertions, 5 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index fbdc1101f36..8ddd3d48539 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -12,8 +12,8 @@ use stable_mir::mir::alloc::GlobalAlloc;
 use stable_mir::mir::mono::{InstanceDef, StaticDef};
 use stable_mir::mir::Body;
 use stable_mir::ty::{
-    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FnDef, GenericArgs, LineInfo,
-    PolyFnSig, RigidTy, Span, TyKind, VariantDef,
+    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FieldDef, FnDef, GenericArgs,
+    LineInfo, PolyFnSig, RigidTy, Span, TyKind, VariantDef,
 };
 use stable_mir::{self, Crate, CrateItem, DefId, Error, Filename, ItemKind, Symbol};
 use std::cell::RefCell;
@@ -219,6 +219,11 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         def.internal(&mut *tables).name.to_string()
     }
 
+    fn variant_fields(&self, def: VariantDef) -> Vec<FieldDef> {
+        let mut tables = self.0.borrow_mut();
+        def.internal(&mut *tables).fields.iter().map(|f| f.stable(&mut *tables)).collect()
+    }
+
     fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error> {
         let mut tables = self.0.borrow_mut();
         let mir_const = cnst.internal(&mut *tables);
@@ -250,6 +255,18 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         tables.tcx.type_of(item.internal(&mut *tables)).instantiate_identity().stable(&mut *tables)
     }
 
+    fn def_ty_with_args(
+        &self,
+        item: stable_mir::DefId,
+        args: &GenericArgs,
+    ) -> Result<stable_mir::ty::Ty, Error> {
+        let mut tables = self.0.borrow_mut();
+        let args = args.internal(&mut *tables);
+        let def_ty = tables.tcx.type_of(item.internal(&mut *tables));
+        // FIXME(celinval): use try_fold instead to avoid crashing.
+        Ok(def_ty.instantiate(tables.tcx, args).stable(&mut *tables))
+    }
+
     fn const_literal(&self, cnst: &stable_mir::ty::Const) -> String {
         internal(cnst).to_string()
     }
diff --git a/compiler/rustc_smir/src/rustc_smir/convert/mod.rs b/compiler/rustc_smir/src/rustc_smir/convert/mod.rs
index 6dca8cfdbb0..c3ee0a60f4d 100644
--- a/compiler/rustc_smir/src/rustc_smir/convert/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/convert/mod.rs
@@ -60,6 +60,14 @@ impl<'tcx> Stable<'tcx> for rustc_hir::CoroutineKind {
     }
 }
 
+impl<'tcx> Stable<'tcx> for rustc_span::Symbol {
+    type T = stable_mir::Symbol;
+
+    fn stable(&self, _tables: &mut Tables<'tcx>) -> Self::T {
+        self.to_string()
+    }
+}
+
 impl<'tcx> Stable<'tcx> for rustc_span::Span {
     type T = stable_mir::ty::Span;
 
@@ -68,6 +76,16 @@ impl<'tcx> Stable<'tcx> for rustc_span::Span {
     }
 }
 
+impl<'tcx, T> Stable<'tcx> for &[T]
+where
+    T: Stable<'tcx>,
+{
+    type T = Vec<T::T>;
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        self.iter().map(|e| e.stable(tables)).collect()
+    }
+}
+
 impl<'tcx, T, U> Stable<'tcx> for (T, U)
 where
     T: Stable<'tcx>,
diff --git a/compiler/rustc_smir/src/rustc_smir/convert/ty.rs b/compiler/rustc_smir/src/rustc_smir/convert/ty.rs
index f837f28e11e..4fe847c291c 100644
--- a/compiler/rustc_smir/src/rustc_smir/convert/ty.rs
+++ b/compiler/rustc_smir/src/rustc_smir/convert/ty.rs
@@ -137,6 +137,17 @@ impl<'tcx> Stable<'tcx> for ty::AdtKind {
     }
 }
 
+impl<'tcx> Stable<'tcx> for ty::FieldDef {
+    type T = stable_mir::ty::FieldDef;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        stable_mir::ty::FieldDef {
+            def: tables.create_def_id(self.did),
+            name: self.name.stable(tables),
+        }
+    }
+}
+
 impl<'tcx> Stable<'tcx> for ty::GenericArgs<'tcx> {
     type T = stable_mir::ty::GenericArgs;
     fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
diff --git a/compiler/stable_mir/src/compiler_interface.rs b/compiler/stable_mir/src/compiler_interface.rs
index 33e301f6973..8f0c5f73796 100644
--- a/compiler/stable_mir/src/compiler_interface.rs
+++ b/compiler/stable_mir/src/compiler_interface.rs
@@ -9,7 +9,7 @@ use crate::mir::alloc::{AllocId, GlobalAlloc};
 use crate::mir::mono::{Instance, InstanceDef, StaticDef};
 use crate::mir::Body;
 use crate::ty::{
-    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FnDef, GenericArgs,
+    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FieldDef, FnDef, GenericArgs,
     GenericPredicates, Generics, ImplDef, ImplTrait, LineInfo, PolyFnSig, RigidTy, Span, TraitDecl,
     TraitDef, Ty, TyKind, VariantDef,
 };
@@ -76,6 +76,7 @@ pub trait Context {
 
     /// The name of a variant.
     fn variant_name(&self, def: VariantDef) -> Symbol;
+    fn variant_fields(&self, def: VariantDef) -> Vec<FieldDef>;
 
     /// Evaluate constant as a target usize.
     fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error>;
@@ -89,6 +90,9 @@ pub trait Context {
     /// Returns the type of given crate item.
     fn def_ty(&self, item: DefId) -> Ty;
 
+    /// Returns the type of given definition instantiated with the given arguments.
+    fn def_ty_with_args(&self, item: DefId, args: &GenericArgs) -> Result<Ty, Error>;
+
     /// Returns literal value of a const as a string.
     fn const_literal(&self, cnst: &Const) -> String;
 
diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
index 3bd42d703c3..6f131cc4f03 100644
--- a/compiler/stable_mir/src/mir/body.rs
+++ b/compiler/stable_mir/src/mir/body.rs
@@ -658,7 +658,7 @@ pub const RETURN_LOCAL: Local = 0;
 /// `b`'s `FieldIdx` is `1`,
 /// `c`'s `FieldIdx` is `0`, and
 /// `g`'s `FieldIdx` is `2`.
-type FieldIdx = usize;
+pub type FieldIdx = usize;
 
 type UserTypeAnnotationIndex = usize;
 
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index ebf43821fb5..2724b1fe0a6 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/compiler/stable_mir/src/ty.rs
@@ -376,6 +376,16 @@ impl AdtDef {
         with(|cx| cx.adt_kind(*self))
     }
 
+    /// Retrieve the type of this Adt.
+    pub fn ty(&self) -> Ty {
+        with(|cx| cx.def_ty(self.0))
+    }
+
+    /// Retrieve the type of this Adt instantiating the type with the given arguments.
+    pub fn ty_with_args(&self, args: &GenericArgs) -> Result<Ty, Error> {
+        with(|cx| cx.def_ty_with_args(self.0, args))
+    }
+
     pub fn is_box(&self) -> bool {
         with(|cx| cx.adt_is_box(*self))
     }
@@ -397,7 +407,7 @@ impl AdtDef {
     }
 
     pub fn variant(&self, idx: VariantIdx) -> Option<VariantDef> {
-        self.variants().get(idx.to_index()).copied()
+        (idx.to_index() < self.num_variants()).then_some(VariantDef { idx, adt_def: *self })
     }
 }
 
@@ -422,6 +432,36 @@ impl VariantDef {
     pub fn name(&self) -> Symbol {
         with(|cx| cx.variant_name(*self))
     }
+
+    /// Retrieve all the fields in this variant.
+    // We expect user to cache this and use it directly since today it is expensive to generate all
+    // fields name.
+    pub fn fields(&self) -> Vec<FieldDef> {
+        with(|cx| cx.variant_fields(*self))
+    }
+}
+
+pub struct FieldDef {
+    /// The field definition.
+    ///
+    /// ## Warning
+    /// Do not access this field directly! This is public for the compiler to have access to it.
+    pub def: DefId,
+
+    /// The field name.
+    pub name: Symbol,
+}
+
+impl FieldDef {
+    /// Retrieve the type of this field instantiating the type with the given arguments.
+    pub fn ty_with_args(&self, args: &GenericArgs) -> Result<Ty, Error> {
+        with(|cx| cx.def_ty_with_args(self.def, args))
+    }
+
+    /// Retrieve the type of this field.
+    pub fn ty(&self) -> Ty {
+        with(|cx| cx.def_ty(self.def))
+    }
 }
 
 impl Display for AdtKind {