Rollup merge of #123659 - celinval:smir-fix-intrinsic, r=oli-obk

Add support to intrinsics fallback body

Before this fix, the call to `body()` would crash, since `has_body()` would return true, but we would try to retrieve the body of an intrinsic which is not allowed.

Instead, the `Instance::body()` function will now convert an Intrinsic into an Item before retrieving its body.

Note: I also changed how we monomorphize the instance body. Unfortunately, the call still ICE for some shims.

r? `@oli-obk`

3 files changed, 152 insertions, 34 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/builder.rs b/compiler/rustc_smir/src/rustc_smir/builder.rs
index 0762016ef75..221224eed01 100644
--- a/compiler/rustc_smir/src/rustc_smir/builder.rs
+++ b/compiler/rustc_smir/src/rustc_smir/builder.rs
@@ -4,9 +4,10 @@
 //! monomorphic body using internal representation.
 //! After that, we convert the internal representation into a stable one.
 use crate::rustc_smir::{Stable, Tables};
+use rustc_hir::def::DefKind;
 use rustc_middle::mir;
 use rustc_middle::mir::visit::MutVisitor;
-use rustc_middle::ty::{self, GenericArgsRef, Ty, TyCtxt};
+use rustc_middle::ty::{self, TyCtxt};
 
 /// Builds a monomorphic body for a given instance.
 pub struct BodyBuilder<'tcx> {
@@ -16,46 +17,43 @@ pub struct BodyBuilder<'tcx> {
 
 impl<'tcx> BodyBuilder<'tcx> {
     pub fn new(tcx: TyCtxt<'tcx>, instance: ty::Instance<'tcx>) -> Self {
+        let instance = match instance.def {
+            // To get the fallback body of an intrinsic, we need to convert it to an item.
+            ty::InstanceDef::Intrinsic(def_id) => ty::Instance::new(def_id, instance.args),
+            _ => instance,
+        };
         BodyBuilder { tcx, instance }
     }
 
     /// Build a stable monomorphic body for a given instance based on the MIR body.
     ///
-    /// Note that we skip instantiation for static and constants. Trying to do so can cause ICE.
-    ///
-    /// We do monomorphize non-generic functions to eval unevaluated constants.
+    /// All constants are also evaluated.
     pub fn build(mut self, tables: &mut Tables<'tcx>) -> stable_mir::mir::Body {
-        let mut body = self.tcx.instance_mir(self.instance.def).clone();
-        if self.tcx.def_kind(self.instance.def_id()).is_fn_like() || !self.instance.args.is_empty()
+        let body = tables.tcx.instance_mir(self.instance.def).clone();
+        let mono_body = if !self.instance.args.is_empty()
+            // Without the `generic_const_exprs` feature gate, anon consts in signatures do not
+            // get generic parameters. Which is wrong, but also not a problem without
+            // generic_const_exprs
+            || self.tcx.def_kind(self.instance.def_id()) != DefKind::AnonConst
         {
-            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),
-        )
+            let mut mono_body = self.instance.instantiate_mir_and_normalize_erasing_regions(
+                tables.tcx,
+                ty::ParamEnv::reveal_all(),
+                ty::EarlyBinder::bind(body),
+            );
+            self.visit_body(&mut mono_body);
+            mono_body
+        } else {
+            // Already monomorphic.
+            body
+        };
+        mono_body.stable(tables)
     }
 }
 
 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 visit_constant(&mut self, constant: &mut mir::ConstOperand<'tcx>, location: mir::Location) {
-        let const_ = self.monomorphize(constant.const_);
+        let const_ = constant.const_;
         let val = match const_.eval(self.tcx, ty::ParamEnv::reveal_all(), constant.span) {
             Ok(v) => v,
             Err(mir::interpret::ErrorHandled::Reported(..)) => return,
@@ -68,10 +66,6 @@ impl<'tcx> MutVisitor<'tcx> for BodyBuilder<'tcx> {
         self.super_constant(constant, location);
     }
 
-    fn visit_args(&mut self, args: &mut GenericArgsRef<'tcx>, _: mir::Location) {
-        *args = self.monomorphize(*args);
-    }
-
     fn tcx(&self) -> TyCtxt<'tcx> {
         self.tcx
     }
diff --git a/compiler/stable_mir/src/mir/mono.rs b/compiler/stable_mir/src/mir/mono.rs
index aafa89c03e0..a032a180fcf 100644
--- a/compiler/stable_mir/src/mir/mono.rs
+++ b/compiler/stable_mir/src/mir/mono.rs
@@ -41,13 +41,22 @@ impl Instance {
         with(|cx| cx.instance_args(self.def))
     }
 
-    /// Get the body of an Instance. The body will be eagerly monomorphized.
+    /// Get the body of an Instance.
+    ///
+    /// The body will be eagerly monomorphized and all constants will already be evaluated.
+    ///
+    /// This method will return the intrinsic fallback body if one was defined.
     pub fn body(&self) -> Option<Body> {
         with(|context| context.instance_body(self.def))
     }
 
     /// Check whether this instance has a body available.
     ///
+    /// For intrinsics with fallback body, this will return `true`. It is up to the user to decide
+    /// whether to specialize the intrinsic or to use its fallback body.
+    ///
+    /// For more information on fallback body, see <https://github.com/rust-lang/rust/issues/93145>.
+    ///
     /// This call is much cheaper than `instance.body().is_some()`, since it doesn't try to build
     /// the StableMIR body.
     pub fn has_body(&self) -> bool {
diff --git a/tests/ui-fulldeps/stable-mir/check_intrinsics.rs b/tests/ui-fulldeps/stable-mir/check_intrinsics.rs
new file mode 100644
index 00000000000..171850b89bb
--- /dev/null
+++ b/tests/ui-fulldeps/stable-mir/check_intrinsics.rs
@@ -0,0 +1,115 @@
+//@ run-pass
+//! Test information regarding intrinsics and ensure we can retrieve the fallback body if it exists.
+//!
+//! This tests relies on the intrinsics implementation, and requires one intrinsic with and one
+//! without a body. It doesn't matter which intrinsic is called here, and feel free to update that
+//! if needed.
+
+//@ ignore-stage1
+//@ ignore-cross-compile
+//@ ignore-remote
+//@ ignore-windows-gnu mingw has troubles with linking https://github.com/rust-lang/rust/pull/116837
+
+#![feature(rustc_private)]
+
+extern crate rustc_hir;
+#[macro_use]
+extern crate rustc_smir;
+extern crate rustc_driver;
+extern crate rustc_interface;
+extern crate stable_mir;
+
+use rustc_smir::rustc_internal;
+use stable_mir::mir::mono::{Instance, InstanceKind};
+use stable_mir::mir::visit::{Location, MirVisitor};
+use stable_mir::mir::{LocalDecl, Terminator, TerminatorKind};
+use stable_mir::ty::{RigidTy, TyKind};
+use std::collections::HashSet;
+use std::convert::TryFrom;
+use std::io::Write;
+use std::ops::ControlFlow;
+
+/// This function tests that we can correctly get type information from binary operations.
+fn test_intrinsics() -> ControlFlow<()> {
+    // Find items in the local crate.
+    let main_def = stable_mir::all_local_items()[0];
+    let main_instance = Instance::try_from(main_def).unwrap();
+    let main_body = main_instance.body().unwrap();
+    let mut visitor = CallsVisitor { locals: main_body.locals(), calls: Default::default() };
+    visitor.visit_body(&main_body);
+
+    let calls = visitor.calls;
+    assert_eq!(calls.len(), 2, "Expected 2 calls, but found: {calls:?}");
+    for intrinsic in &calls {
+        check_intrinsic(intrinsic)
+    }
+
+    ControlFlow::Continue(())
+}
+
+/// This check is unfortunately tight to the implementation of intrinsics.
+///
+/// We want to ensure that StableMIR can handle intrinsics with and without fallback body.
+///
+/// If by any chance this test breaks because you changed how an intrinsic is implemented, please
+/// update the test to invoke a different intrinsic.
+fn check_intrinsic(intrinsic: &Instance) {
+    assert_eq!(intrinsic.kind, InstanceKind::Intrinsic);
+    let name = intrinsic.intrinsic_name().unwrap();
+    if intrinsic.has_body() {
+        let Some(body) = intrinsic.body() else { unreachable!("Expected a body") };
+        assert!(!body.blocks.is_empty());
+        assert_eq!(&name, "likely");
+    } else {
+        assert!(intrinsic.body().is_none());
+        assert_eq!(&name, "size_of_val");
+    }
+}
+
+struct CallsVisitor<'a> {
+    locals: &'a [LocalDecl],
+    calls: HashSet<Instance>,
+}
+
+impl<'a> MirVisitor for CallsVisitor<'a> {
+    fn visit_terminator(&mut self, term: &Terminator, _loc: Location) {
+        match &term.kind {
+            TerminatorKind::Call { func, .. } => {
+                let TyKind::RigidTy(RigidTy::FnDef(def, args)) =
+                    func.ty(self.locals).unwrap().kind()
+                    else {
+                        return;
+                    };
+                self.calls.insert(Instance::resolve(def, &args).unwrap());
+            }
+            _ => {}
+        }
+    }
+}
+
+/// 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 = "binop_input.rs";
+    generate_input(&path).unwrap();
+    let args = vec!["rustc".to_string(), "--crate-type=lib".to_string(), path.to_string()];
+    run!(args, test_intrinsics).unwrap();
+}
+
+fn generate_input(path: &str) -> std::io::Result<()> {
+    let mut file = std::fs::File::create(path)?;
+    write!(
+        file,
+        r#"
+        #![feature(core_intrinsics)]
+        use std::intrinsics::*;
+        pub fn use_intrinsics(init: bool) -> bool {{
+            let sz = unsafe {{ size_of_val("hi") }};
+            likely(init && sz == 2)
+        }}
+        "#
+    )?;
+    Ok(())
+}