Rollup merge of #118524 - celinval:smir-instance-def, r=ouz-a

Add more information to StableMIR Instance

Allow stable MIR users to retrieve an instance function signature, the index for a VTable instance and more information about its underlying definition.

These are needed to properly interpret function calls, either via VTable or direct calls. The `CrateDef` implementation will also allow users to emit diagnostic messages.

I also fixed a few issues that we had identified before with how we were retrieving body of things that may not have a body available.

compiler/rustc_smir/src/rustc_smir/context.rs

@@ -3,6 +3,7 @@
 //! This trait is currently the main interface between the Rust compiler,
 //! and the `stable_mir` crate.
 
+use rustc_middle::ty;
 use rustc_middle::ty::print::{with_forced_trimmed_paths, with_no_trimmed_paths};
 use rustc_middle::ty::{GenericPredicates, Instance, ParamEnv, ScalarInt, ValTree};
 use rustc_span::def_id::LOCAL_CRATE;
@@ -12,9 +13,9 @@ 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,
-    RigidTy, Span, TyKind,
+    PolyFnSig, RigidTy, Span, TyKind,
 };
-use stable_mir::{self, Crate, CrateItem, Error, Filename, ItemKind, Symbol};
+use stable_mir::{self, Crate, CrateItem, DefId, Error, Filename, ItemKind, Symbol};
 use std::cell::RefCell;
 
 use crate::rustc_internal::{internal, RustcInternal};
@@ -39,6 +40,12 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         tables.tcx.instance_mir(rustc_middle::ty::InstanceDef::Item(def_id)).stable(&mut tables)
     }
 
+    fn has_body(&self, def: DefId) -> bool {
+        let tables = self.0.borrow();
+        let def_id = tables[def];
+        tables.tcx.is_mir_available(def_id)
+    }
+
     fn all_trait_decls(&self) -> stable_mir::TraitDecls {
         let mut tables = self.0.borrow_mut();
         tables
@@ -195,6 +202,13 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         def.internal(&mut *tables).is_box()
     }
 
+    fn fn_sig(&self, def: FnDef, args: &GenericArgs) -> PolyFnSig {
+        let mut tables = self.0.borrow_mut();
+        let def_id = def.0.internal(&mut *tables);
+        let sig = tables.tcx.fn_sig(def_id).instantiate(tables.tcx, args.internal(&mut *tables));
+        sig.stable(&mut *tables)
+    }
+
     fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error> {
         let mut tables = self.0.borrow_mut();
         let mir_const = cnst.internal(&mut *tables);
@@ -266,6 +280,12 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         tables.tcx.symbol_name(instance).name.to_string()
     }
 
+    fn is_empty_drop_shim(&self, def: InstanceDef) -> bool {
+        let tables = self.0.borrow_mut();
+        let instance = tables.instances[def];
+        matches!(instance.def, ty::InstanceDef::DropGlue(_, None))
+    }
+
     fn mono_instance(&self, item: stable_mir::CrateItem) -> stable_mir::mir::mono::Instance {
         let mut tables = self.0.borrow_mut();
         let def_id = tables[item.0];

compiler/rustc_smir/src/rustc_smir/convert/ty.rs

@@ -777,7 +777,9 @@ impl<'tcx> Stable<'tcx> for ty::Instance<'tcx> {
         let kind = match self.def {
             ty::InstanceDef::Item(..) => stable_mir::mir::mono::InstanceKind::Item,
             ty::InstanceDef::Intrinsic(..) => stable_mir::mir::mono::InstanceKind::Intrinsic,
-            ty::InstanceDef::Virtual(..) => stable_mir::mir::mono::InstanceKind::Virtual,
+            ty::InstanceDef::Virtual(_def_id, idx) => {
+                stable_mir::mir::mono::InstanceKind::Virtual { idx }
+            }
             ty::InstanceDef::VTableShim(..)
             | ty::InstanceDef::ReifyShim(..)
             | ty::InstanceDef::FnPtrAddrShim(..)

compiler/stable_mir/src/compiler_interface.rs

@@ -10,8 +10,8 @@ use crate::mir::mono::{Instance, InstanceDef, StaticDef};
 use crate::mir::Body;
 use crate::ty::{
     AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, FnDef, GenericArgs,
-    GenericPredicates, Generics, ImplDef, ImplTrait, LineInfo, RigidTy, Span, TraitDecl, TraitDef,
-    Ty, TyKind,
+    GenericPredicates, Generics, ImplDef, ImplTrait, LineInfo, PolyFnSig, RigidTy, Span, TraitDecl,
+    TraitDef, Ty, TyKind,
 };
 use crate::{
     mir, Crate, CrateItem, CrateItems, DefId, Error, Filename, ImplTraitDecls, ItemKind, Symbol,
@@ -24,7 +24,11 @@ pub trait Context {
     fn entry_fn(&self) -> Option<CrateItem>;
     /// Retrieve all items of the local crate that have a MIR associated with them.
     fn all_local_items(&self) -> CrateItems;
+    /// Retrieve the body of a function.
+    /// This function will panic if the body is not available.
     fn mir_body(&self, item: DefId) -> mir::Body;
+    /// Check whether the body of a function is available.
+    fn has_body(&self, item: DefId) -> bool;
     fn all_trait_decls(&self) -> TraitDecls;
     fn trait_decl(&self, trait_def: &TraitDef) -> TraitDecl;
     fn all_trait_impls(&self) -> ImplTraitDecls;
@@ -64,6 +68,9 @@ pub trait Context {
     /// Returns if the ADT is a box.
     fn adt_is_box(&self, def: AdtDef) -> bool;
 
+    /// Retrieve the function signature for the given generic arguments.
+    fn fn_sig(&self, def: FnDef, args: &GenericArgs) -> PolyFnSig;
+
     /// Evaluate constant as a target usize.
     fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error>;
 
@@ -85,8 +92,7 @@ pub trait Context {
     /// Obtain the representation of a type.
     fn ty_kind(&self, ty: Ty) -> TyKind;
 
-    /// Get the body of an Instance.
-    /// FIXME: Monomorphize the body.
+    /// Get the body of an Instance which is already monomorphized.
     fn instance_body(&self, instance: InstanceDef) -> Option<Body>;
 
     /// Get the instance type with generic substitutions applied and lifetimes erased.
@@ -98,6 +104,9 @@ pub trait Context {
     /// Get the instance mangled name.
     fn instance_mangled_name(&self, instance: InstanceDef) -> Symbol;
 
+    /// Check if this is an empty DropGlue shim.
+    fn is_empty_drop_shim(&self, def: InstanceDef) -> bool;
+
     /// Convert a non-generic crate item into an instance.
     /// This function will panic if the item is generic.
     fn mono_instance(&self, item: CrateItem) -> Instance;

compiler/stable_mir/src/mir/mono.rs

@@ -1,6 +1,6 @@
 use crate::crate_def::CrateDef;
 use crate::mir::Body;
-use crate::ty::{Allocation, ClosureDef, ClosureKind, FnDef, GenericArgs, IndexedVal, Ty};
+use crate::ty::{Allocation, ClosureDef, ClosureKind, FnDef, FnSig, GenericArgs, IndexedVal, Ty};
 use crate::{with, CrateItem, DefId, Error, ItemKind, Opaque, Symbol};
 use std::fmt::{Debug, Formatter};
 
@@ -27,7 +27,8 @@ pub enum InstanceKind {
     /// A compiler intrinsic function.
     Intrinsic,
     /// A virtual function definition stored in a VTable.
-    Virtual,
+    /// The `idx` field indicates the position in the VTable for this instance.
+    Virtual { idx: usize },
     /// A compiler generated shim.
     Shim,
 }
@@ -106,6 +107,24 @@ impl Instance {
             })
         })
     }
+
+    /// Get this function signature with all types already instantiated.
+    pub fn fn_sig(&self) -> FnSig {
+        self.ty().kind().fn_sig().unwrap().skip_binder()
+    }
+
+    /// Check whether this instance is an empty shim.
+    ///
+    /// Allow users to check if this shim can be ignored when called directly.
+    ///
+    /// We have decided not to export different types of Shims to StableMIR users, however, this
+    /// is a query that can be very helpful for users when processing DropGlue.
+    ///
+    /// When generating code for a Drop terminator, users can ignore an empty drop glue.
+    /// These shims are only needed to generate a valid Drop call done via VTable.
+    pub fn is_empty_shim(&self) -> bool {
+        self.kind == InstanceKind::Shim && with(|cx| cx.is_empty_drop_shim(self.def))
+    }
 }
 
 impl Debug for Instance {
@@ -124,8 +143,6 @@ impl TryFrom<CrateItem> for Instance {
 
     fn try_from(item: CrateItem) -> Result<Self, Self::Error> {
         with(|context| {
-            // FIXME(celinval):
-            // - Return `Err` if instance does not have a body.
             if !context.requires_monomorphization(item.0) {
                 Ok(context.mono_instance(item))
             } else {
@@ -141,11 +158,13 @@ impl TryFrom<Instance> for CrateItem {
     type Error = crate::Error;
 
     fn try_from(value: Instance) -> Result<Self, Self::Error> {
-        if value.kind == InstanceKind::Item {
-            Ok(CrateItem(with(|context| context.instance_def_id(value.def))))
-        } else {
-            Err(Error::new(format!("Item kind `{:?}` cannot be converted", value.kind)))
-        }
+        with(|context| {
+            if value.kind == InstanceKind::Item && context.has_body(value.def.def_id()) {
+                Ok(CrateItem(context.instance_def_id(value.def)))
+            } else {
+                Err(Error::new(format!("Item kind `{:?}` cannot be converted", value.kind)))
+            }
+        })
     }
 }
 
@@ -170,6 +189,12 @@ impl From<StaticDef> for CrateItem {
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct InstanceDef(usize);
 
+impl CrateDef for InstanceDef {
+    fn def_id(&self) -> DefId {
+        with(|context| context.instance_def_id(*self))
+    }
+}
+
 crate_def! {
     /// Holds information about a static variable definition.
     pub StaticDef;

compiler/stable_mir/src/ty.rs

@@ -168,7 +168,11 @@ impl TyKind {
     }
 
     pub fn is_unit(&self) -> bool {
-        matches!(self, TyKind::RigidTy(RigidTy::Tuple(data)) if data.len() == 0)
+        matches!(self, TyKind::RigidTy(RigidTy::Tuple(data)) if data.is_empty())
+    }
+
+    pub fn is_bool(&self) -> bool {
+        matches!(self, TyKind::RigidTy(RigidTy::Bool))
     }
 
     pub fn is_trait(&self) -> bool {
@@ -187,6 +191,14 @@ impl TyKind {
         matches!(self, TyKind::RigidTy(RigidTy::Adt(def, _)) if def.kind() == AdtKind::Union)
     }
 
+    pub fn is_fn(&self) -> bool {
+        matches!(self, TyKind::RigidTy(RigidTy::FnDef(..)))
+    }
+
+    pub fn is_fn_ptr(&self) -> bool {
+        matches!(self, TyKind::RigidTy(RigidTy::FnPtr(..)))
+    }
+
     pub fn trait_principal(&self) -> Option<Binder<ExistentialTraitRef>> {
         if let TyKind::RigidTy(RigidTy::Dynamic(predicates, _, _)) = self {
             if let Some(Binder { value: ExistentialPredicate::Trait(trait_ref), bound_vars }) =
@@ -227,6 +239,15 @@ impl TyKind {
             _ => None,
         }
     }
+
+    /// Get the function signature for function like types (Fn, FnPtr, Closure, Coroutine)
+    pub fn fn_sig(&self) -> Option<PolyFnSig> {
+        match self {
+            TyKind::RigidTy(RigidTy::FnDef(def, args)) => Some(with(|cx| cx.fn_sig(*def, args))),
+            TyKind::RigidTy(RigidTy::FnPtr(sig)) => Some(sig.clone()),
+            _ => None,
+        }
+    }
 }
 
 pub struct TypeAndMut {
@@ -307,8 +328,9 @@ crate_def! {
 }
 
 impl FnDef {
-    pub fn body(&self) -> Body {
-        with(|ctx| ctx.mir_body(self.0))
+    // Get the function body if available.
+    pub fn body(&self) -> Option<Body> {
+        with(|ctx| ctx.has_body(self.0).then(|| ctx.mir_body(self.0)))
     }
 }
 
@@ -488,6 +510,16 @@ pub struct FnSig {
     pub abi: Abi,
 }
 
+impl FnSig {
+    pub fn output(&self) -> Ty {
+        self.inputs_and_output[self.inputs_and_output.len() - 1]
+    }
+
+    pub fn inputs(&self) -> &[Ty] {
+        &self.inputs_and_output[..self.inputs_and_output.len() - 1]
+    }
+}
+
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub enum Abi {
     Rust,

tests/ui-fulldeps/stable-mir/check_defs.rs

@@ -18,10 +18,11 @@ extern crate rustc_driver;
 extern crate rustc_interface;
 extern crate stable_mir;
 
+use std::assert_matches::assert_matches;
 use mir::{mono::Instance, TerminatorKind::*};
 use rustc_middle::ty::TyCtxt;
 use rustc_smir::rustc_internal;
-use stable_mir::ty::{RigidTy, TyKind};
+use stable_mir::ty::{RigidTy, TyKind, Ty, UintTy};
 use stable_mir::*;
 use std::io::Write;
 use std::ops::ControlFlow;
@@ -39,6 +40,7 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> {
     assert_eq!(instances.len(), 2);
     test_fn(instances[0], "from_u32", "std::char::from_u32", "core");
     test_fn(instances[1], "Vec::<u8>::new", "std::vec::Vec::<u8>::new", "alloc");
+    test_vec_new(instances[1]);
     ControlFlow::Continue(())
 }
 
@@ -56,6 +58,30 @@ fn test_fn(instance: Instance, expected_trimmed: &str, expected_qualified: &str,
     assert_eq!(&item.krate().name, krate);
 }
 
+fn extract_elem_ty(ty: Ty) -> Ty {
+    match ty.kind() {
+        TyKind::RigidTy(RigidTy::Adt(_, args)) => {
+            *args.0[0].expect_ty()
+        }
+        _ => unreachable!("Expected Vec ADT, but found: {ty:?}")
+    }
+}
+
+/// Check signature and type of `Vec::<u8>::new` and its generic version.
+fn test_vec_new(instance: mir::mono::Instance) {
+    let sig = instance.fn_sig();
+    assert_matches!(sig.inputs(), &[]);
+    let elem_ty = extract_elem_ty(sig.output());
+    assert_matches!(elem_ty.kind(), TyKind::RigidTy(RigidTy::Uint(UintTy::U8)));
+
+    // Get the signature for Vec::<T>::new.
+    let item = CrateItem::try_from(instance).unwrap();
+    let ty = item.ty();
+    let gen_sig = ty.kind().fn_sig().unwrap().skip_binder();
+    let gen_ty = extract_elem_ty(gen_sig.output());
+    assert_matches!(gen_ty.kind(), TyKind::Param(_));
+}
+
 /// Inspect the instance body
 fn get_instances(body: mir::Body) -> Vec<Instance> {
     body.blocks.iter().filter_map(|bb| {