Add Variant and a few more APIs to stable_mir

compiler/rustc_smir/src/rustc_internal/internal.rs

@@ -9,11 +9,7 @@ use rustc_middle::ty::{self as rustc_ty, Ty as InternalTy};
 use rustc_span::Symbol;
 use stable_mir::mir::alloc::AllocId;
 use stable_mir::mir::mono::{Instance, MonoItem, StaticDef};
-use stable_mir::ty::{
-    AdtDef, Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const,
-    ExistentialTraitRef, FloatTy, GenericArgKind, GenericArgs, IntTy, Region, RigidTy, Span,
-    TraitRef, Ty, UintTy,
-};
+use stable_mir::ty::{AdtDef, Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const, ExistentialTraitRef, FloatTy, GenericArgKind, GenericArgs, IndexedVal, IntTy, Region, RigidTy, Span, TraitRef, Ty, UintTy, VariantDef, VariantIdx};
 use stable_mir::{CrateItem, DefId};
 
 use super::RustcInternal;
@@ -141,6 +137,22 @@ impl<'tcx> RustcInternal<'tcx> for FloatTy {
     }
 }
 
+impl<'tcx> RustcInternal<'tcx> for VariantIdx {
+    type T = rustc_target::abi::VariantIdx;
+
+    fn internal(&self, _tables: &mut Tables<'tcx>) -> Self::T {
+        rustc_target::abi::VariantIdx::from(self.to_index())
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for VariantDef {
+    type T = &'tcx rustc_ty::VariantDef;
+
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        self.adt_def.internal(tables).variant(self.idx.internal(tables))
+    }
+}
+
 fn ty_const<'tcx>(constant: &Const, tables: &mut Tables<'tcx>) -> rustc_ty::Const<'tcx> {
     match constant.internal(tables) {
         rustc_middle::mir::Const::Ty(c) => c,

compiler/rustc_smir/src/rustc_smir/context.rs

@@ -13,7 +13,7 @@ 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,
+    PolyFnSig, RigidTy, Span, TyKind, VariantDef,
 };
 use stable_mir::{self, Crate, CrateItem, DefId, Error, Filename, ItemKind, Symbol};
 use std::cell::RefCell;
@@ -209,6 +209,16 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         sig.stable(&mut *tables)
     }
 
+    fn adt_variants_len(&self, def: AdtDef) -> usize {
+        let mut tables = self.0.borrow_mut();
+        def.internal(&mut *tables).variants().len()
+    }
+
+    fn variant_name(&self, def: VariantDef) -> Symbol {
+        let mut tables = self.0.borrow_mut();
+        def.internal(&mut *tables).name.to_string()
+    }
+
     fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error> {
         let mut tables = self.0.borrow_mut();
         let mir_const = cnst.internal(&mut *tables);

compiler/rustc_smir/src/rustc_smir/convert/mir.rs

@@ -517,7 +517,7 @@ impl<'tcx> Stable<'tcx> for mir::AggregateKind<'tcx> {
             mir::AggregateKind::Adt(def_id, var_idx, generic_arg, user_ty_index, field_idx) => {
                 stable_mir::mir::AggregateKind::Adt(
                     tables.adt_def(*def_id),
-                    var_idx.index(),
+                    var_idx.stable(tables),
                     generic_arg.stable(tables),
                     user_ty_index.map(|idx| idx.index()),
                     field_idx.map(|idx| idx.index()),

compiler/rustc_smir/src/rustc_smir/convert/mod.rs

@@ -1,7 +1,7 @@
 //! Conversion of internal Rust compiler items to stable ones.
 
 use rustc_target::abi::FieldIdx;
-use stable_mir::mir::VariantIdx;
+use stable_mir::ty::{IndexedVal, VariantIdx};
 
 use crate::rustc_smir::{Stable, Tables};
 
@@ -25,17 +25,10 @@ impl<'tcx> Stable<'tcx> for FieldIdx {
     }
 }
 
-impl<'tcx> Stable<'tcx> for (rustc_target::abi::VariantIdx, FieldIdx) {
-    type T = (usize, usize);
-    fn stable(&self, _: &mut Tables<'tcx>) -> Self::T {
-        (self.0.as_usize(), self.1.as_usize())
-    }
-}
-
 impl<'tcx> Stable<'tcx> for rustc_target::abi::VariantIdx {
     type T = VariantIdx;
     fn stable(&self, _: &mut Tables<'tcx>) -> Self::T {
-        self.as_usize()
+        VariantIdx::to_val(self.as_usize())
     }
 }
 
@@ -74,3 +67,14 @@ impl<'tcx> Stable<'tcx> for rustc_span::Span {
         tables.create_span(*self)
     }
 }
+
+impl<'tcx, T, U> Stable<'tcx> for (T, U)
+where
+    T: Stable<'tcx>,
+    U: Stable<'tcx>,
+{
+    type T = (T::T, U::T);
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        (self.0.stable(tables), self.1.stable(tables))
+    }
+}

compiler/stable_mir/src/compiler_interface.rs

@@ -11,7 +11,7 @@ use crate::mir::Body;
 use crate::ty::{
     AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FnDef, GenericArgs,
     GenericPredicates, Generics, ImplDef, ImplTrait, LineInfo, PolyFnSig, RigidTy, Span, TraitDecl,
-    TraitDef, Ty, TyKind,
+    TraitDef, Ty, TyKind, VariantDef,
 };
 use crate::{
     mir, Crate, CrateItem, CrateItems, DefId, Error, Filename, ImplTraitDecls, ItemKind, Symbol,
@@ -71,6 +71,12 @@ pub trait Context {
     /// Retrieve the function signature for the given generic arguments.
     fn fn_sig(&self, def: FnDef, args: &GenericArgs) -> PolyFnSig;
 
+    /// The number of variants in this ADT.
+    fn adt_variants_len(&self, def: AdtDef) -> usize;
+
+    /// The name of a variant.
+    fn variant_name(&self, def: VariantDef) -> Symbol;
+
     /// Evaluate constant as a target usize.
     fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error>;
 

compiler/stable_mir/src/mir/body.rs

@@ -1,6 +1,7 @@
 use crate::mir::pretty::{function_body, pretty_statement, pretty_terminator};
 use crate::ty::{
     AdtDef, ClosureDef, Const, CoroutineDef, GenericArgs, Movability, Region, RigidTy, Ty, TyKind,
+    VariantIdx,
 };
 use crate::{Error, Opaque, Span, Symbol};
 use std::io;
@@ -9,17 +10,17 @@ use std::io;
 pub struct Body {
     pub blocks: Vec<BasicBlock>,
 
-    // Declarations of locals within the function.
-    //
-    // The first local is the return value pointer, followed by `arg_count`
-    // locals for the function arguments, followed by any user-declared
-    // variables and temporaries.
+    /// Declarations of locals within the function.
+    ///
+    /// The first local is the return value pointer, followed by `arg_count`
+    /// locals for the function arguments, followed by any user-declared
+    /// variables and temporaries.
     pub(super) locals: LocalDecls,
 
-    // The number of arguments this function takes.
+    /// The number of arguments this function takes.
     pub(super) arg_count: usize,
 
-    // Debug information pertaining to user variables, including captures.
+    /// Debug information pertaining to user variables, including captures.
     pub(super) var_debug_info: Vec<VarDebugInfo>,
 }
 
@@ -69,6 +70,11 @@ impl Body {
         &self.locals
     }
 
+    /// Get the local declaration for this local.
+    pub fn local_decl(&self, local: Local) -> Option<&LocalDecl> {
+        self.locals.get(local)
+    }
+
     pub fn dump<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
         writeln!(w, "{}", function_body(self))?;
         self.blocks
@@ -492,12 +498,32 @@ pub struct Place {
     pub projection: Vec<ProjectionElem>,
 }
 
+impl From<Local> for Place {
+    fn from(local: Local) -> Self {
+        Place { local, projection: vec![] }
+    }
+}
+
+/// Debug information pertaining to a user variable.
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct VarDebugInfo {
+    /// The variable name.
     pub name: Symbol,
+
+    /// Source info of the user variable, including the scope
+    /// within which the variable is visible (to debuginfo)
     pub source_info: SourceInfo,
+
+    /// The user variable's data is split across several fragments,
+    /// each described by a `VarDebugInfoFragment`.
     pub composite: Option<VarDebugInfoFragment>,
+
+    /// Where the data for this user variable is to be found.
     pub value: VarDebugInfoContents,
+
+    /// When present, indicates what argument number this variable is in the function that it
+    /// originated from (starting from 1). Note, if MIR inlining is enabled, then this is the
+    /// argument number in the original function before it was inlined.
     pub argument_index: Option<u16>,
 }
 
@@ -634,21 +660,6 @@ pub const RETURN_LOCAL: Local = 0;
 /// `g`'s `FieldIdx` is `2`.
 type FieldIdx = usize;
 
-/// The source-order index of a variant in a type.
-///
-/// For example, in the following types,
-/// ```ignore(illustrative)
-/// enum Demo1 {
-///    Variant0 { a: bool, b: i32 },
-///    Variant1 { c: u8, d: u64 },
-/// }
-/// struct Demo2 { e: u8, f: u16, g: u8 }
-/// ```
-/// `a` is in the variant with the `VariantIdx` of `0`,
-/// `c` is in the variant with the `VariantIdx` of `1`, and
-/// `g` is in the variant with the `VariantIdx` of `0`.
-pub type VariantIdx = usize;
-
 type UserTypeAnnotationIndex = usize;
 
 #[derive(Clone, Debug, Eq, PartialEq)]

compiler/stable_mir/src/ty.rs

@@ -30,6 +30,16 @@ impl Ty {
     pub fn try_new_array(elem_ty: Ty, size: u64) -> Result<Ty, Error> {
         Ok(Ty::from_rigid_kind(RigidTy::Array(elem_ty, Const::try_from_target_usize(size)?)))
     }
+
+    /// Create a new pointer type.
+    pub fn new_ptr(pointee_ty: Ty, mutability: Mutability) -> Ty {
+        Ty::from_rigid_kind(RigidTy::RawPtr(pointee_ty, mutability))
+    }
+
+    /// Create a type representing `usize`.
+    pub fn usize_ty() -> Ty {
+        Ty::from_rigid_kind(RigidTy::Uint(UintTy::Usize))
+    }
 }
 
 impl Ty {
@@ -369,6 +379,49 @@ impl AdtDef {
     pub fn is_box(&self) -> bool {
         with(|cx| cx.adt_is_box(*self))
     }
+
+    /// The number of variants in this ADT.
+    pub fn num_variants(&self) -> usize {
+        with(|cx| cx.adt_variants_len(*self))
+    }
+
+    /// Retrieve the variants in this ADT.
+    pub fn variants(&self) -> Vec<VariantDef> {
+        self.variants_iter().collect()
+    }
+
+    /// Iterate over the variants in this ADT.
+    pub fn variants_iter(&self) -> impl Iterator<Item = VariantDef> + '_ {
+        (0..self.num_variants())
+            .map(|idx| VariantDef { idx: VariantIdx::to_val(idx), adt_def: *self })
+    }
+
+    pub fn variant(&self, idx: VariantIdx) -> Option<VariantDef> {
+        self.variants().get(idx.to_index()).copied()
+    }
+}
+
+/// Definition of a variant, which can be either a struct / union field or an enum variant.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct VariantDef {
+    /// The variant index.
+    ///
+    /// ## Warning
+    /// Do not access this field directly!
+    pub idx: VariantIdx,
+    /// The data type where this variant comes from.
+    /// For now, we use this to retrieve information about the variant itself so we don't need to
+    /// cache more information.
+    ///
+    /// ## Warning
+    /// Do not access this field directly!
+    pub adt_def: AdtDef,
+}
+
+impl VariantDef {
+    pub fn name(&self) -> Symbol {
+        with(|cx| cx.variant_name(*self))
+    }
 }
 
 impl Display for AdtKind {
@@ -906,3 +959,21 @@ macro_rules! index_impl {
 index_impl!(ConstId);
 index_impl!(Ty);
 index_impl!(Span);
+
+/// The source-order index of a variant in a type.
+///
+/// For example, in the following types,
+/// ```ignore(illustrative)
+/// enum Demo1 {
+///    Variant0 { a: bool, b: i32 },
+///    Variant1 { c: u8, d: u64 },
+/// }
+/// struct Demo2 { e: u8, f: u16, g: u8 }
+/// ```
+/// `a` is in the variant with the `VariantIdx` of `0`,
+/// `c` is in the variant with the `VariantIdx` of `1`, and
+/// `g` is in the variant with the `VariantIdx` of `0`.
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct VariantIdx(usize);
+
+index_impl!(VariantIdx);