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

Add stable Instance::body() and RustcInternal trait

The `Instance::body()` returns a monomorphized body.

For that, we had to implement visitor that monomorphize types and constants. We are also introducing the RustcInternal trait that will allow us to convert back from Stable to Internal.

Note that this trait is not yet visible for our users as it depends on Tables. We should probably add a new trait that can be exposed.

The tests here are very simple, and I'm planning on creating more exhaustive tests in the project-mir repo. But I was hoping to get some feedback here first.

r? ```@oli-obk```

compiler/rustc_smir/src/rustc_internal/internal.rs

@@ -0,0 +1,61 @@
+//! Module containing the translation from stable mir constructs to the rustc counterpart.
+//!
+//! This module will only include a few constructs to allow users to invoke internal rustc APIs
+//! due to incomplete stable coverage.
+
+// Prefer importing stable_mir over internal rustc constructs to make this file more readable.
+use crate::rustc_smir::{MaybeStable, Tables};
+use rustc_middle::ty::{self as rustc_ty, Ty as InternalTy};
+use stable_mir::ty::{Const, GenericArgKind, GenericArgs, Region, Ty};
+use stable_mir::DefId;
+
+use super::RustcInternal;
+
+impl<'tcx> RustcInternal<'tcx> for DefId {
+    type T = rustc_span::def_id::DefId;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.def_ids[*self]
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for GenericArgs {
+    type T = rustc_ty::GenericArgsRef<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.tcx.mk_args_from_iter(self.0.iter().map(|arg| arg.internal(tables)))
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for GenericArgKind {
+    type T = rustc_ty::GenericArg<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        match self {
+            GenericArgKind::Lifetime(reg) => reg.internal(tables).into(),
+            GenericArgKind::Type(ty) => ty.internal(tables).into(),
+            GenericArgKind::Const(cnst) => cnst.internal(tables).into(),
+        }
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for Region {
+    type T = rustc_ty::Region<'tcx>;
+    fn internal(&self, _tables: &mut Tables<'tcx>) -> Self::T {
+        todo!()
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for Ty {
+    type T = InternalTy<'tcx>;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        match tables.types[self.0] {
+            MaybeStable::Stable(_) => todo!(),
+            MaybeStable::Rustc(ty) => ty,
+        }
+    }
+}
+
+impl<'tcx> RustcInternal<'tcx> for Const {
+    type T = rustc_ty::Const<'tcx>;
+    fn internal(&self, _tables: &mut Tables<'tcx>) -> Self::T {
+        todo!()
+    }
+}

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -20,6 +20,8 @@ use std::fmt::Debug;
 use std::hash::Hash;
 use std::ops::{ControlFlow, Index};
 
+mod internal;
+
 impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
     type Output = DefId;
 
@@ -231,3 +233,11 @@ impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> Index<V
         k
     }
 }
+
+/// Trait used to translate a stable construct to its rustc counterpart.
+///
+/// This is basically a mirror of [crate::rustc_smir::Stable].
+pub(crate) trait RustcInternal<'tcx> {
+    type T;
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T;
+}

compiler/rustc_smir/src/rustc_smir/builder.rs

@@ -0,0 +1,55 @@
+//! Logic required to produce a monomorphic stable body.
+//!
+//! We first retrieve and monomorphize the rustc body representation, i.e., we generate a
+//! monomorphic body using internal representation.
+//! After that, we convert the internal representation into a stable one.
+use crate::rustc_smir::{Stable, Tables};
+use rustc_middle::mir;
+use rustc_middle::mir::visit::MutVisitor;
+use rustc_middle::ty::{self, Ty, TyCtxt};
+
+/// Builds a monomorphic body for a given instance.
+pub struct BodyBuilder<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    instance: ty::Instance<'tcx>,
+}
+
+impl<'tcx> BodyBuilder<'tcx> {
+    pub fn new(tcx: TyCtxt<'tcx>, instance: ty::Instance<'tcx>) -> Self {
+        BodyBuilder { tcx, instance }
+    }
+
+    pub fn build(mut self, tables: &mut Tables<'tcx>) -> stable_mir::mir::Body {
+        let mut body = self.tcx.instance_mir(self.instance.def).clone();
+        let generics = self.tcx.generics_of(self.instance.def_id());
+        if generics.requires_monomorphization(self.tcx) {
+            self.visit_body(&mut body);
+        }
+        body.stable(tables)
+    }
+
+    fn monomorphize<T>(&self, value: T) -> T
+    where
+        T: ty::TypeFoldable<TyCtxt<'tcx>>,
+    {
+        self.instance.instantiate_mir_and_normalize_erasing_regions(
+            self.tcx,
+            ty::ParamEnv::reveal_all(),
+            ty::EarlyBinder::bind(value),
+        )
+    }
+}
+
+impl<'tcx> MutVisitor<'tcx> for BodyBuilder<'tcx> {
+    fn visit_ty_const(&mut self, ct: &mut ty::Const<'tcx>, _location: mir::Location) {
+        *ct = self.monomorphize(*ct);
+    }
+
+    fn visit_ty(&mut self, ty: &mut Ty<'tcx>, _: mir::visit::TyContext) {
+        *ty = self.monomorphize(*ty);
+    }
+
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+}

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -7,7 +7,7 @@
 //!
 //! For now, we are developing everything inside `rustc`, thus, we keep this module private.
 
-use crate::rustc_internal::IndexMap;
+use crate::rustc_internal::{IndexMap, RustcInternal};
 use crate::rustc_smir::hir::def::DefKind;
 use crate::rustc_smir::stable_mir::ty::{BoundRegion, EarlyBoundRegion, Region};
 use rustc_hir as hir;
@@ -26,6 +26,7 @@ use stable_mir::{self, opaque, Context, Filename};
 use tracing::debug;
 
 mod alloc;
+mod builder;
 
 impl<'tcx> Context for Tables<'tcx> {
     fn local_crate(&self) -> stable_mir::Crate {
@@ -171,8 +172,9 @@ impl<'tcx> Context for Tables<'tcx> {
         }
     }
 
-    fn instance_body(&mut self, _def: InstanceDef) -> Body {
-        todo!("Monomorphize the body")
+    fn instance_body(&mut self, def: InstanceDef) -> Body {
+        let instance = self.instances[def];
+        builder::BodyBuilder::new(self.tcx, instance).build(self)
     }
 
     fn instance_ty(&mut self, def: InstanceDef) -> stable_mir::ty::Ty {
@@ -195,9 +197,21 @@ impl<'tcx> Context for Tables<'tcx> {
         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
     }
+
+    fn resolve_instance(
+        &mut self,
+        def: stable_mir::ty::FnDef,
+        args: &stable_mir::ty::GenericArgs,
+    ) -> Option<stable_mir::mir::mono::Instance> {
+        let def_id = def.0.internal(self);
+        let args_ref = args.internal(self);
+        match Instance::resolve(self.tcx, ParamEnv::reveal_all(), def_id, args_ref) {
+            Ok(Some(instance)) => Some(instance.stable(self)),
+            Ok(None) | Err(_) => None,
+        }
+    }
 }
 
 #[derive(Clone)]

compiler/stable_mir/src/error.rs

@@ -4,6 +4,7 @@
 //! - [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::convert::From;
 use std::fmt::{Debug, Display, Formatter};
 use std::{error, fmt};
 
@@ -31,6 +32,12 @@ impl Error {
     }
 }
 
+impl From<&str> for Error {
+    fn from(value: &str) -> Self {
+        Self(value.into())
+    }
+}
+
 impl Display for Error {
     fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
         Display::fmt(&self.0, f)

compiler/stable_mir/src/lib.rs

@@ -39,6 +39,7 @@ pub mod visitor;
 
 pub use error::*;
 use mir::mono::Instance;
+use ty::{FnDef, GenericArgs};
 
 /// Use String for now but we should replace it.
 pub type Symbol = String;
@@ -233,6 +234,9 @@ pub trait Context {
 
     /// Item requires monomorphization.
     fn requires_monomorphization(&self, def_id: DefId) -> bool;
+
+    /// Resolve an instance from the given function definition and generic arguments.
+    fn resolve_instance(&mut self, def: FnDef, args: &GenericArgs) -> Option<Instance>;
 }
 
 // A thread local variable that stores a pointer to the tables mapping between TyCtxt

compiler/stable_mir/src/mir/body.rs

@@ -5,9 +5,11 @@ use crate::{ty::Ty, Span};
 #[derive(Clone, Debug)]
 pub struct Body {
     pub blocks: Vec<BasicBlock>,
-    pub locals: Vec<LocalDecl>,
+    pub locals: LocalDecls,
 }
 
+type LocalDecls = Vec<LocalDecl>;
+
 #[derive(Clone, Debug)]
 pub struct LocalDecl {
     pub ty: Ty,
@@ -344,6 +346,7 @@ pub enum Operand {
 #[derive(Clone, Debug)]
 pub struct Place {
     pub local: Local,
+    /// projection out of a place (access a field, deref a pointer, etc)
     pub projection: String,
 }
 
@@ -462,3 +465,25 @@ pub enum NullOp {
     /// Returns the offset of a field.
     OffsetOf(Vec<FieldIdx>),
 }
+
+impl Operand {
+    pub fn ty(&self, locals: &LocalDecls) -> Ty {
+        match self {
+            Operand::Copy(place) | Operand::Move(place) => place.ty(locals),
+            Operand::Constant(c) => c.ty(),
+        }
+    }
+}
+
+impl Constant {
+    pub fn ty(&self) -> Ty {
+        self.literal.ty
+    }
+}
+
+impl Place {
+    pub fn ty(&self, locals: &LocalDecls) -> Ty {
+        let _start_ty = locals[self.local].ty;
+        todo!("Implement projection")
+    }
+}

compiler/stable_mir/src/mir/mono.rs

@@ -1,5 +1,5 @@
 use crate::mir::Body;
-use crate::ty::{IndexedVal, Ty};
+use crate::ty::{FnDef, GenericArgs, IndexedVal, Ty};
 use crate::{with, CrateItem, DefId, Error, Opaque};
 use std::fmt::Debug;
 
@@ -41,6 +41,15 @@ impl Instance {
     pub fn ty(&self) -> Ty {
         with(|context| context.instance_ty(self.def))
     }
+
+    /// Resolve an instance starting from a function definition and generic arguments.
+    pub fn resolve(def: FnDef, args: &GenericArgs) -> Result<Instance, crate::Error> {
+        with(|context| {
+            context.resolve_instance(def, args).ok_or_else(|| {
+                crate::Error::new(format!("Failed to resolve `{def:?}` with `{args:?}`"))
+            })
+        })
+    }
 }
 
 /// Try to convert a crate item into an instance.

compiler/stable_mir/src/ty.rs

@@ -225,6 +225,7 @@ pub struct ImplDef(pub DefId);
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct RegionDef(pub DefId);
 
+/// A list of generic arguments.
 #[derive(Clone, Debug)]
 pub struct GenericArgs(pub Vec<GenericArgKind>);
 

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

@@ -15,7 +15,8 @@ extern crate rustc_smir;
 extern crate stable_mir;
 
 use rustc_middle::ty::TyCtxt;
-
+use mir::{mono::Instance, TerminatorKind::*};
+use stable_mir::ty::{TyKind, RigidTy};
 use stable_mir::*;
 use rustc_smir::rustc_internal;
 use std::io::Write;
@@ -43,9 +44,28 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> {
     // For all generic items, try_from should fail.
     assert!(generic.iter().all(|item| mir::mono::Instance::try_from(*item).is_err()));
 
+    for instance in instances {
+        test_body(instance.body())
+    }
     ControlFlow::Continue(())
 }
 
+/// Inspect the instance body
+fn test_body(body: mir::Body) {
+    for term in body.blocks.iter().map(|bb| &bb.terminator) {
+        match &term.kind {
+            Call{ func, .. } => {
+                let TyKind::RigidTy(ty) = func.ty(&body.locals).kind() else { unreachable!() };
+                let RigidTy::FnDef(def, args) = ty else { unreachable!() };
+                let result = Instance::resolve(def, &args);
+                assert!(result.is_ok());
+            }
+            Goto {..} | Assert{..} | SwitchInt{..} | Return | Drop {..} => { /* Do nothing */}
+            _ => { unreachable!("Unexpected terminator {term:?}") }
+        }
+    }
+}
+
 
 /// 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.
@@ -56,6 +76,7 @@ fn main() {
     generate_input(&path).unwrap();
     let args = vec![
         "rustc".to_string(),
+        "-Cpanic=abort".to_string(),
         "--crate-type=lib".to_string(),
         "--crate-name".to_string(),
         CRATE_NAME.to_string(),
@@ -78,6 +99,9 @@ fn generate_input(path: &str) -> std::io::Result<()> {
     }}
 
     pub fn monomorphic() {{
+        let v = vec![10];
+        let dup = ty_param(&v);
+        assert_eq!(v, dup);
     }}
 
     pub mod foo {{