Rollup merge of #116719 - celinval:smir-mono, r=oli-obk

Add MonoItems and Instance to stable_mir

Also add a few methods to instantiate instances and get an instance definition. We're still missing support to actually monomorphize the instance body.

This is related to https://github.com/rust-lang/project-stable-mir/issues/36

r? ``@oli-obk``

``@oli-obk`` is that what you were thinking? I incorporated ``@bjorn3`` idea of just adding a Shim instance definition in https://github.com/rust-lang/rust/pull/116465.

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -6,9 +6,11 @@
 use crate::rustc_internal;
 use crate::rustc_smir::Tables;
 use rustc_data_structures::fx;
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_driver::{Callbacks, Compilation, RunCompiler};
 use rustc_interface::{interface, Queries};
 use rustc_middle::mir::interpret::AllocId;
+use rustc_middle::ty;
 use rustc_middle::ty::TyCtxt;
 use rustc_span::def_id::{CrateNum, DefId};
 use rustc_span::Span;
@@ -97,7 +99,7 @@ impl<'tcx> Tables<'tcx> {
         stable_mir::ty::Prov(self.create_alloc_id(aid))
     }
 
-    fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
+    pub(crate) fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
         self.def_ids.create_or_fetch(did)
     }
 
@@ -108,6 +110,17 @@ impl<'tcx> Tables<'tcx> {
     pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span {
         self.spans.create_or_fetch(span)
     }
+
+    pub(crate) fn instance_def(
+        &mut self,
+        instance: ty::Instance<'tcx>,
+    ) -> stable_mir::mir::mono::InstanceDef {
+        self.instances.create_or_fetch(instance)
+    }
+
+    pub(crate) fn static_def(&mut self, did: DefId) -> stable_mir::mir::mono::StaticDef {
+        stable_mir::mir::mono::StaticDef(self.create_def_id(did))
+    }
 }
 
 pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
@@ -118,10 +131,11 @@ pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
     stable_mir::run(
         Tables {
             tcx,
-            def_ids: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() },
-            alloc_ids: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() },
-            spans: rustc_internal::IndexMap { index_map: fx::FxIndexMap::default() },
+            def_ids: IndexMap::default(),
+            alloc_ids: IndexMap::default(),
+            spans: IndexMap::default(),
             types: vec![],
+            instances: IndexMap::default(),
         },
         f,
     );
@@ -192,6 +206,12 @@ pub struct IndexMap<K, V> {
     index_map: fx::FxIndexMap<K, V>,
 }
 
+impl<K, V> Default for IndexMap<K, V> {
+    fn default() -> Self {
+        Self { index_map: FxIndexMap::default() }
+    }
+}
+
 impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> IndexMap<K, V> {
     pub fn create_or_fetch(&mut self, key: K) -> V {
         let len = self.index_map.len();

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -13,10 +13,12 @@ use crate::rustc_smir::stable_mir::ty::{BoundRegion, EarlyBoundRegion, Region};
 use rustc_hir as hir;
 use rustc_middle::mir;
 use rustc_middle::mir::interpret::{alloc_range, AllocId};
-use rustc_middle::ty::{self, Ty, TyCtxt, Variance};
+use rustc_middle::mir::mono::MonoItem;
+use rustc_middle::ty::{self, Instance, ParamEnv, Ty, TyCtxt, Variance};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use rustc_target::abi::FieldIdx;
-use stable_mir::mir::{CopyNonOverlapping, Statement, UserTypeProjection, VariantIdx};
+use stable_mir::mir::mono::InstanceDef;
+use stable_mir::mir::{Body, CopyNonOverlapping, Statement, UserTypeProjection, VariantIdx};
 use stable_mir::ty::{
     FloatTy, GenericParamDef, IntTy, LineInfo, Movability, RigidTy, Span, TyKind, UintTy,
 };
@@ -119,29 +121,7 @@ impl<'tcx> Context for Tables<'tcx> {
 
     fn mir_body(&mut self, item: stable_mir::DefId) -> stable_mir::mir::Body {
         let def_id = self[item];
-        let mir = self.tcx.instance_mir(ty::InstanceDef::Item(def_id));
-        stable_mir::mir::Body {
-            blocks: mir
-                .basic_blocks
-                .iter()
-                .map(|block| stable_mir::mir::BasicBlock {
-                    terminator: block.terminator().stable(self),
-                    statements: block
-                        .statements
-                        .iter()
-                        .map(|statement| statement.stable(self))
-                        .collect(),
-                })
-                .collect(),
-            locals: mir
-                .local_decls
-                .iter()
-                .map(|decl| stable_mir::mir::LocalDecl {
-                    ty: self.intern_ty(decl.ty),
-                    span: decl.source_info.span.stable(self),
-                })
-                .collect(),
-        }
+        self.tcx.instance_mir(ty::InstanceDef::Item(def_id)).stable(self)
     }
 
     fn ty_kind(&mut self, ty: stable_mir::ty::Ty) -> TyKind {
@@ -190,6 +170,34 @@ impl<'tcx> Context for Tables<'tcx> {
                 .collect(),
         }
     }
+
+    fn instance_body(&mut self, _def: InstanceDef) -> Body {
+        todo!("Monomorphize the body")
+    }
+
+    fn instance_ty(&mut self, def: InstanceDef) -> stable_mir::ty::Ty {
+        let instance = self.instances[def];
+        let ty = instance.ty(self.tcx, ParamEnv::empty());
+        self.intern_ty(ty)
+    }
+
+    fn instance_def_id(&mut self, def: InstanceDef) -> stable_mir::DefId {
+        let def_id = self.instances[def].def_id();
+        self.create_def_id(def_id)
+    }
+
+    fn mono_instance(&mut self, item: stable_mir::CrateItem) -> stable_mir::mir::mono::Instance {
+        let def_id = self[item.0];
+        Instance::mono(self.tcx, def_id).stable(self)
+    }
+
+    fn requires_monomorphization(&self, def_id: stable_mir::DefId) -> bool {
+        let def_id = self[def_id];
+        let generics = self.tcx.generics_of(def_id);
+        let result = generics.requires_monomorphization(self.tcx);
+        println!("req {result}: {def_id:?}");
+        result
+    }
 }
 
 #[derive(Clone)]
@@ -224,7 +232,8 @@ pub struct Tables<'tcx> {
     pub def_ids: IndexMap<DefId, stable_mir::DefId>,
     pub alloc_ids: IndexMap<AllocId, stable_mir::AllocId>,
     pub spans: IndexMap<rustc_span::Span, Span>,
-    pub types: Vec<MaybeStable<stable_mir::ty::TyKind, Ty<'tcx>>>,
+    pub types: Vec<MaybeStable<TyKind, Ty<'tcx>>>,
+    pub instances: IndexMap<ty::Instance<'tcx>, InstanceDef>,
 }
 
 impl<'tcx> Tables<'tcx> {
@@ -254,6 +263,35 @@ pub(crate) trait Stable<'tcx> {
     fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T;
 }
 
+impl<'tcx> Stable<'tcx> for mir::Body<'tcx> {
+    type T = stable_mir::mir::Body;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        stable_mir::mir::Body {
+            blocks: self
+                .basic_blocks
+                .iter()
+                .map(|block| stable_mir::mir::BasicBlock {
+                    terminator: block.terminator().stable(tables),
+                    statements: block
+                        .statements
+                        .iter()
+                        .map(|statement| statement.stable(tables))
+                        .collect(),
+                })
+                .collect(),
+            locals: self
+                .local_decls
+                .iter()
+                .map(|decl| stable_mir::mir::LocalDecl {
+                    ty: tables.intern_ty(decl.ty),
+                    span: decl.source_info.span.stable(tables),
+                })
+                .collect(),
+        }
+    }
+}
+
 impl<'tcx> Stable<'tcx> for mir::Statement<'tcx> {
     type T = stable_mir::mir::Statement;
     fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
@@ -1637,3 +1675,38 @@ impl<'tcx> Stable<'tcx> for DefKind {
         opaque(self)
     }
 }
+
+impl<'tcx> Stable<'tcx> for ty::Instance<'tcx> {
+    type T = stable_mir::mir::mono::Instance;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        let def = tables.instance_def(*self);
+        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::VTableShim(..)
+            | ty::InstanceDef::ReifyShim(..)
+            | ty::InstanceDef::FnPtrAddrShim(..)
+            | ty::InstanceDef::ClosureOnceShim { .. }
+            | ty::InstanceDef::ThreadLocalShim(..)
+            | ty::InstanceDef::DropGlue(..)
+            | ty::InstanceDef::CloneShim(..)
+            | ty::InstanceDef::FnPtrShim(..) => stable_mir::mir::mono::InstanceKind::Shim,
+        };
+        stable_mir::mir::mono::Instance { def, kind }
+    }
+}
+
+impl<'tcx> Stable<'tcx> for MonoItem<'tcx> {
+    type T = stable_mir::mir::mono::MonoItem;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        use stable_mir::mir::mono::MonoItem as StableMonoItem;
+        match self {
+            MonoItem::Fn(instance) => StableMonoItem::Fn(instance.stable(tables)),
+            MonoItem::Static(def_id) => StableMonoItem::Static(tables.static_def(*def_id)),
+            MonoItem::GlobalAsm(item_id) => StableMonoItem::GlobalAsm(opaque(item_id)),
+        }
+    }
+}

compiler/stable_mir/src/error.rs

@@ -0,0 +1,69 @@
+//! When things go wrong, we need some error handling.
+//! There are a few different types of errors in StableMIR:
+//!
+//! - [CompilerError]: This represents errors that can be raised when invoking the compiler.
+//! - [Error]: Generic error that represents the reason why a request that could not be fulfilled.
+
+use std::fmt::{Debug, Display, Formatter};
+use std::{error, fmt};
+
+/// An error type used to represent an error that has already been reported by the compiler.
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum CompilerError<T> {
+    /// Internal compiler error (I.e.: Compiler crashed).
+    ICE,
+    /// Compilation failed.
+    CompilationFailed,
+    /// Compilation was interrupted.
+    Interrupted(T),
+    /// Compilation skipped. This happens when users invoke rustc to retrieve information such as
+    /// --version.
+    Skipped,
+}
+
+/// A generic error to represent an API request that cannot be fulfilled.
+#[derive(Debug)]
+pub struct Error(String);
+
+impl Error {
+    pub(crate) fn new(msg: String) -> Self {
+        Self(msg)
+    }
+}
+
+impl Display for Error {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        Display::fmt(&self.0, f)
+    }
+}
+
+impl<T> Display for CompilerError<T>
+where
+    T: Display,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            CompilerError::ICE => write!(f, "Internal Compiler Error"),
+            CompilerError::CompilationFailed => write!(f, "Compilation Failed"),
+            CompilerError::Interrupted(reason) => write!(f, "Compilation Interrupted: {reason}"),
+            CompilerError::Skipped => write!(f, "Compilation Skipped"),
+        }
+    }
+}
+
+impl<T> Debug for CompilerError<T>
+where
+    T: Debug,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            CompilerError::ICE => write!(f, "Internal Compiler Error"),
+            CompilerError::CompilationFailed => write!(f, "Compilation Failed"),
+            CompilerError::Interrupted(reason) => write!(f, "Compilation Interrupted: {reason:?}"),
+            CompilerError::Skipped => write!(f, "Compilation Skipped"),
+        }
+    }
+}
+
+impl error::Error for Error {}
+impl<T> error::Error for CompilerError<T> where T: Display + Debug {}

compiler/stable_mir/src/lib.rs

@@ -17,6 +17,8 @@
 //! The goal is to eventually be published on
 //! [crates.io](https://crates.io).
 
+use crate::mir::mono::InstanceDef;
+use crate::mir::Body;
 use std::cell::Cell;
 use std::fmt;
 use std::fmt::Debug;
@@ -29,11 +31,15 @@ use self::ty::{
 #[macro_use]
 extern crate scoped_tls;
 
+pub mod error;
 pub mod fold;
 pub mod mir;
 pub mod ty;
 pub mod visitor;
 
+pub use error::*;
+use mir::mono::Instance;
+
 /// Use String for now but we should replace it.
 pub type Symbol = String;
 
@@ -85,20 +91,6 @@ pub type TraitDecls = Vec<TraitDef>;
 /// A list of impl trait decls.
 pub type ImplTraitDecls = Vec<ImplDef>;
 
-/// An error type used to represent an error that has already been reported by the compiler.
-#[derive(Clone, Copy, Debug, PartialEq, Eq)]
-pub enum CompilerError<T> {
-    /// Internal compiler error (I.e.: Compiler crashed).
-    ICE,
-    /// Compilation failed.
-    CompilationFailed,
-    /// Compilation was interrupted.
-    Interrupted(T),
-    /// Compilation skipped. This happens when users invoke rustc to retrieve information such as
-    /// --version.
-    Skipped,
-}
-
 /// Holds information about a crate.
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Crate {
@@ -113,7 +105,7 @@ pub type Filename = Opaque;
 /// Holds information about an item in the crate.
 /// For now, it only stores the item DefId. Use functions inside `rustc_internal` module to
 /// use this item.
-#[derive(Clone, PartialEq, Eq, Debug)]
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
 pub struct CrateItem(pub DefId);
 
 impl CrateItem {
@@ -132,6 +124,10 @@ impl CrateItem {
     pub fn kind(&self) -> DefKind {
         with(|cx| cx.def_kind(self.0))
     }
+
+    pub fn requires_monomorphization(&self) -> bool {
+        with(|cx| cx.requires_monomorphization(self.0))
+    }
 }
 
 /// Return the function where execution starts if the current
@@ -220,6 +216,23 @@ pub trait Context {
 
     /// Create a new `Ty` from scratch without information from rustc.
     fn mk_ty(&mut self, kind: TyKind) -> Ty;
+
+    /// Get the body of an Instance.
+    /// FIXME: Monomorphize the body.
+    fn instance_body(&mut self, instance: InstanceDef) -> Body;
+
+    /// Get the instance type with generic substitutions applied and lifetimes erased.
+    fn instance_ty(&mut self, instance: InstanceDef) -> Ty;
+
+    /// Get the instance.
+    fn instance_def_id(&mut self, instance: InstanceDef) -> DefId;
+
+    /// Convert a non-generic crate item into an instance.
+    /// This function will panic if the item is generic.
+    fn mono_instance(&mut self, item: CrateItem) -> Instance;
+
+    /// Item requires monomorphization.
+    fn requires_monomorphization(&self, def_id: DefId) -> bool;
 }
 
 // A thread local variable that stores a pointer to the tables mapping between TyCtxt

compiler/stable_mir/src/mir.rs

@@ -1,3 +1,4 @@
 mod body;
+pub mod mono;
 
 pub use body::*;

compiler/stable_mir/src/mir/mono.rs

@@ -0,0 +1,89 @@
+use crate::mir::Body;
+use crate::ty::{IndexedVal, Ty};
+use crate::{with, CrateItem, DefId, Error, Opaque};
+use std::fmt::Debug;
+
+#[derive(Clone, Debug)]
+pub enum MonoItem {
+    Fn(Instance),
+    Static(StaticDef),
+    GlobalAsm(Opaque),
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct Instance {
+    /// The type of instance.
+    pub kind: InstanceKind,
+    /// An ID used to get the instance definition from the compiler.
+    /// Do not use this field directly.
+    pub def: InstanceDef,
+}
+
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum InstanceKind {
+    /// A user defined item.
+    Item,
+    /// A compiler intrinsic function.
+    Intrinsic,
+    /// A virtual function definition stored in a VTable.
+    Virtual,
+    /// A compiler generated shim.
+    Shim,
+}
+
+impl Instance {
+    /// Get the body of an Instance. The body will be eagerly monomorphized.
+    pub fn body(&self) -> Body {
+        with(|context| context.instance_body(self.def))
+    }
+
+    /// Get the instance type with generic substitutions applied and lifetimes erased.
+    pub fn ty(&self) -> Ty {
+        with(|context| context.instance_ty(self.def))
+    }
+}
+
+/// Try to convert a crate item into an instance.
+/// The item cannot be generic in order to be converted into an instance.
+impl TryFrom<CrateItem> for Instance {
+    type Error = crate::Error;
+
+    fn try_from(item: CrateItem) -> Result<Self, Self::Error> {
+        with(|context| {
+            if !context.requires_monomorphization(item.0) {
+                Ok(context.mono_instance(item))
+            } else {
+                Err(Error::new("Item requires monomorphization".to_string()))
+            }
+        })
+    }
+}
+
+/// Try to convert an instance into a crate item.
+/// Only user defined instances can be converted.
+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)))
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct InstanceDef(usize);
+
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub struct StaticDef(pub DefId);
+
+impl IndexedVal for InstanceDef {
+    fn to_val(index: usize) -> Self {
+        InstanceDef(index)
+    }
+    fn to_index(&self) -> usize {
+        self.0
+    }
+}

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

@@ -0,0 +1,91 @@
+// run-pass
+// Test that users are able to use stable mir APIs to retrieve monomorphized instances
+
+// ignore-stage1
+// ignore-cross-compile
+// ignore-remote
+// edition: 2021
+
+#![feature(rustc_private)]
+#![feature(assert_matches)]
+#![feature(control_flow_enum)]
+
+extern crate rustc_middle;
+extern crate rustc_smir;
+extern crate stable_mir;
+
+use rustc_middle::ty::TyCtxt;
+
+use stable_mir::*;
+use rustc_smir::rustc_internal;
+use std::io::Write;
+use std::ops::ControlFlow;
+
+const CRATE_NAME: &str = "input";
+
+/// This function uses the Stable MIR APIs to get information about the test crate.
+fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> {
+    let items = stable_mir::all_local_items();
+
+    // Get all items and split generic vs monomorphic items.
+    let (generic, mono) : (Vec<_>, Vec<_>) = items.into_iter().partition(|item| {
+        item.requires_monomorphization()
+    });
+    assert_eq!(mono.len(), 3, "Expected 2 mono functions and one constant");
+    assert_eq!(generic.len(), 2, "Expected 2 generic functions");
+
+    // For all monomorphic items, get the correspondent instances.
+    let instances = mono.iter().filter_map(|item| {
+        mir::mono::Instance::try_from(*item).ok()
+    }).collect::<Vec<mir::mono::Instance>>();
+    assert_eq!(instances.len(), mono.len());
+
+    // For all generic items, try_from should fail.
+    assert!(generic.iter().all(|item| mir::mono::Instance::try_from(*item).is_err()));
+
+    ControlFlow::Continue(())
+}
+
+
+/// This test will generate and analyze a dummy crate using the stable mir.
+/// For that, it will first write the dummy crate into a file.
+/// Then it will create a `StableMir` using custom arguments and then
+/// it will run the compiler.
+fn main() {
+    let path = "instance_input.rs";
+    generate_input(&path).unwrap();
+    let args = vec![
+        "rustc".to_string(),
+        "--crate-type=lib".to_string(),
+        "--crate-name".to_string(),
+        CRATE_NAME.to_string(),
+        path.to_string(),
+    ];
+    rustc_internal::StableMir::new(args, test_stable_mir).run().unwrap();
+}
+
+fn generate_input(path: &str) -> std::io::Result<()> {
+    let mut file = std::fs::File::create(path)?;
+    write!(
+        file,
+        r#"
+    pub fn ty_param<T>(t: &T) -> T where T: Clone {{
+        t.clone()
+    }}
+
+    pub fn const_param<const LEN: usize>(a: [bool; LEN]) -> bool {{
+        LEN > 0 && a[0]
+    }}
+
+    pub fn monomorphic() {{
+    }}
+
+    pub mod foo {{
+        pub fn bar_mono(i: i32) -> i64 {{
+            i as i64
+        }}
+    }}
+    "#
+    )?;
+    Ok(())
+}