Add Variant and a few more APIs to stable_mir

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