introduce a macro for shim signature checking

Co-Authored-By: geetanshjuneja <ronitjuneja2002@gmail.com>

Author: RalfJung

src/tools/miri/src/helpers.rs

@@ -3,7 +3,7 @@ use std::time::Duration;
 use std::{cmp, iter};
 
 use rand::RngCore;
-use rustc_abi::{Align, CanonAbi, ExternAbi, FieldIdx, FieldsShape, Size, Variants};
+use rustc_abi::{Align, ExternAbi, FieldIdx, FieldsShape, Size, Variants};
 use rustc_apfloat::Float;
 use rustc_apfloat::ieee::{Double, Half, Quad, Single};
 use rustc_hir::Safety;
@@ -14,11 +14,10 @@ use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::middle::dependency_format::Linkage;
 use rustc_middle::middle::exported_symbols::ExportedSymbol;
 use rustc_middle::ty::layout::{LayoutOf, MaybeResult, TyAndLayout};
-use rustc_middle::ty::{self, Binder, FloatTy, FnSig, IntTy, Ty, TyCtxt, UintTy};
+use rustc_middle::ty::{self, FloatTy, IntTy, Ty, TyCtxt, UintTy};
 use rustc_session::config::CrateType;
 use rustc_span::{Span, Symbol};
 use rustc_symbol_mangling::mangle_internal_symbol;
-use rustc_target::callconv::FnAbi;
 
 use crate::*;
 
@@ -437,7 +436,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     /// For now, arguments must be scalars (so that the caller does not have to know the layout).
     ///
     /// If you do not provide a return place, a dangling zero-sized place will be created
-    /// for your convenience.
+    /// for your convenience. This is only valid if the return type is `()`.
     fn call_function(
         &mut self,
         f: ty::Instance<'tcx>,
@@ -452,7 +451,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let mir = this.load_mir(f.def, None)?;
         let dest = match dest {
             Some(dest) => dest.clone(),
-            None => MPlaceTy::fake_alloc_zst(this.layout_of(mir.return_ty())?),
+            None => MPlaceTy::fake_alloc_zst(this.machine.layouts.unit),
         };
 
         // Construct a function pointer type representing the caller perspective.
@@ -465,6 +464,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         );
         let caller_fn_abi = this.fn_abi_of_fn_ptr(ty::Binder::dummy(sig), ty::List::empty())?;
 
+        // This will also show proper errors if there is any ABI mismatch.
         this.init_stack_frame(
             f,
             mir,
@@ -929,21 +929,6 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         self.read_c_str_with_char_size(ptr, wchar_t.size, wchar_t.align.abi)
     }
 
-    /// Check that the calling convention is what we expect.
-    fn check_callconv<'a>(
-        &self,
-        fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
-        exp_abi: CanonAbi,
-    ) -> InterpResult<'a, ()> {
-        if fn_abi.conv != exp_abi {
-            throw_ub_format!(
-                r#"calling a function with calling convention "{exp_abi}" using caller calling convention "{}""#,
-                fn_abi.conv
-            );
-        }
-        interp_ok(())
-    }
-
     fn frame_in_std(&self) -> bool {
         let this = self.eval_context_ref();
         let frame = this.frame();
@@ -967,161 +952,6 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         crate_name == "std" || crate_name == "std_miri_test"
     }
 
-    fn check_abi_and_shim_symbol_clash(
-        &mut self,
-        abi: &FnAbi<'tcx, Ty<'tcx>>,
-        exp_abi: CanonAbi,
-        link_name: Symbol,
-    ) -> InterpResult<'tcx, ()> {
-        self.check_callconv(abi, exp_abi)?;
-        if let Some((body, instance)) = self.eval_context_mut().lookup_exported_symbol(link_name)? {
-            // If compiler-builtins is providing the symbol, then don't treat it as a clash.
-            // We'll use our built-in implementation in `emulate_foreign_item_inner` for increased
-            // performance. Note that this means we won't catch any undefined behavior in
-            // compiler-builtins when running other crates, but Miri can still be run on
-            // compiler-builtins itself (or any crate that uses it as a normal dependency)
-            if self.eval_context_ref().tcx.is_compiler_builtins(instance.def_id().krate) {
-                return interp_ok(());
-            }
-
-            throw_machine_stop!(TerminationInfo::SymbolShimClashing {
-                link_name,
-                span: body.span.data(),
-            })
-        }
-        interp_ok(())
-    }
-
-    fn check_shim<'a, const N: usize>(
-        &mut self,
-        abi: &FnAbi<'tcx, Ty<'tcx>>,
-        exp_abi: CanonAbi,
-        link_name: Symbol,
-        args: &'a [OpTy<'tcx>],
-    ) -> InterpResult<'tcx, &'a [OpTy<'tcx>; N]> {
-        self.check_abi_and_shim_symbol_clash(abi, exp_abi, link_name)?;
-
-        if abi.c_variadic {
-            throw_ub_format!(
-                "calling a non-variadic function with a variadic caller-side signature"
-            );
-        }
-        if let Ok(ops) = args.try_into() {
-            return interp_ok(ops);
-        }
-        throw_ub_format!(
-            "incorrect number of arguments for `{link_name}`: got {}, expected {}",
-            args.len(),
-            N
-        )
-    }
-
-    /// Check that the given `caller_fn_abi` matches the expected ABI described by
-    /// `callee_abi`, `callee_input_tys`, `callee_output_ty`, and then returns the list of
-    /// arguments.
-    fn check_shim_abi<'a, const N: usize>(
-        &mut self,
-        link_name: Symbol,
-        caller_fn_abi: &FnAbi<'tcx, Ty<'tcx>>,
-        callee_abi: ExternAbi,
-        callee_input_tys: [Ty<'tcx>; N],
-        callee_output_ty: Ty<'tcx>,
-        caller_args: &'a [OpTy<'tcx>],
-    ) -> InterpResult<'tcx, &'a [OpTy<'tcx>; N]> {
-        let this = self.eval_context_mut();
-        let mut inputs_and_output = callee_input_tys.to_vec();
-        inputs_and_output.push(callee_output_ty);
-        let fn_sig_binder = Binder::dummy(FnSig {
-            inputs_and_output: this.machine.tcx.mk_type_list(&inputs_and_output),
-            c_variadic: false,
-            // This does not matter for the ABI.
-            safety: Safety::Safe,
-            abi: callee_abi,
-        });
-        let callee_fn_abi = this.fn_abi_of_fn_ptr(fn_sig_binder, Default::default())?;
-
-        this.check_abi_and_shim_symbol_clash(caller_fn_abi, callee_fn_abi.conv, link_name)?;
-
-        if caller_fn_abi.c_variadic {
-            throw_ub_format!(
-                "ABI mismatch: calling a non-variadic function with a variadic caller-side signature"
-            );
-        }
-
-        if callee_fn_abi.fixed_count != caller_fn_abi.fixed_count {
-            throw_ub_format!(
-                "ABI mismatch: expected {} arguments, found {} arguments ",
-                callee_fn_abi.fixed_count,
-                caller_fn_abi.fixed_count
-            );
-        }
-
-        if callee_fn_abi.can_unwind && !caller_fn_abi.can_unwind {
-            throw_ub_format!(
-                "ABI mismatch: callee may unwind, but caller-side signature prohibits unwinding",
-            );
-        }
-
-        if !this.check_argument_compat(&caller_fn_abi.ret, &callee_fn_abi.ret)? {
-            throw_ub!(AbiMismatchReturn {
-                caller_ty: caller_fn_abi.ret.layout.ty,
-                callee_ty: callee_fn_abi.ret.layout.ty
-            });
-        }
-
-        if let Some(index) = caller_fn_abi
-            .args
-            .iter()
-            .zip(callee_fn_abi.args.iter())
-            .map(|(caller_arg, callee_arg)| this.check_argument_compat(caller_arg, callee_arg))
-            .collect::<InterpResult<'tcx, Vec<bool>>>()?
-            .into_iter()
-            .position(|b| !b)
-        {
-            throw_ub!(AbiMismatchArgument {
-                caller_ty: caller_fn_abi.args[index].layout.ty,
-                callee_ty: callee_fn_abi.args[index].layout.ty
-            });
-        }
-
-        if let Ok(ops) = caller_args.try_into() {
-            return interp_ok(ops);
-        }
-        unreachable!()
-    }
-
-    /// Check shim for variadic function.
-    /// Returns a tuple that consisting of an array of fixed args, and a slice of varargs.
-    fn check_shim_variadic<'a, const N: usize>(
-        &mut self,
-        abi: &FnAbi<'tcx, Ty<'tcx>>,
-        exp_abi: CanonAbi,
-        link_name: Symbol,
-        args: &'a [OpTy<'tcx>],
-    ) -> InterpResult<'tcx, (&'a [OpTy<'tcx>; N], &'a [OpTy<'tcx>])>
-    where
-        &'a [OpTy<'tcx>; N]: TryFrom<&'a [OpTy<'tcx>]>,
-    {
-        self.check_abi_and_shim_symbol_clash(abi, exp_abi, link_name)?;
-
-        if !abi.c_variadic {
-            throw_ub_format!(
-                "calling a variadic function with a non-variadic caller-side signature"
-            );
-        }
-        if abi.fixed_count != u32::try_from(N).unwrap() {
-            throw_ub_format!(
-                "incorrect number of fixed arguments for variadic function `{}`: got {}, expected {N}",
-                link_name.as_str(),
-                abi.fixed_count
-            )
-        }
-        if let Some(args) = args.split_first_chunk() {
-            return interp_ok(args);
-        }
-        panic!("mismatch between signature and `args` slice");
-    }
-
     /// Mark a machine allocation that was just created as immutable.
     fn mark_immutable(&mut self, mplace: &MPlaceTy<'tcx>) {
         let this = self.eval_context_mut();
@@ -1317,39 +1147,6 @@ impl<'tcx> MiriMachine<'tcx> {
     }
 }
 
-/// Check that the number of args is what we expect.
-pub fn check_intrinsic_arg_count<'a, 'tcx, const N: usize>(
-    args: &'a [OpTy<'tcx>],
-) -> InterpResult<'tcx, &'a [OpTy<'tcx>; N]>
-where
-    &'a [OpTy<'tcx>; N]: TryFrom<&'a [OpTy<'tcx>]>,
-{
-    if let Ok(ops) = args.try_into() {
-        return interp_ok(ops);
-    }
-    throw_ub_format!(
-        "incorrect number of arguments for intrinsic: got {}, expected {}",
-        args.len(),
-        N
-    )
-}
-
-/// Check that the number of varargs is at least the minimum what we expect.
-/// Fixed args should not be included.
-pub fn check_min_vararg_count<'a, 'tcx, const N: usize>(
-    name: &'a str,
-    args: &'a [OpTy<'tcx>],
-) -> InterpResult<'tcx, &'a [OpTy<'tcx>; N]> {
-    if let Some((ops, _)) = args.split_first_chunk() {
-        return interp_ok(ops);
-    }
-    throw_ub_format!(
-        "not enough variadic arguments for `{name}`: got {}, expected at least {}",
-        args.len(),
-        N
-    )
-}
-
 pub fn isolation_abort_error<'tcx>(name: &str) -> InterpResult<'tcx> {
     throw_machine_stop!(TerminationInfo::UnsupportedInIsolation(format!(
         "{name} not available when isolation is enabled",

src/tools/miri/src/intrinsics/atomic.rs

@@ -2,7 +2,7 @@ use rustc_middle::mir::BinOp;
 use rustc_middle::ty::AtomicOrdering;
 use rustc_middle::{mir, ty};
 
-use self::helpers::check_intrinsic_arg_count;
+use super::check_intrinsic_arg_count;
 use crate::*;
 
 pub enum AtomicOp {

src/tools/miri/src/intrinsics/mod.rs

@@ -14,11 +14,28 @@ use rustc_middle::ty::{self, FloatTy, ScalarInt};
 use rustc_span::{Symbol, sym};
 
 use self::atomic::EvalContextExt as _;
-use self::helpers::{ToHost, ToSoft, check_intrinsic_arg_count};
+use self::helpers::{ToHost, ToSoft};
 use self::simd::EvalContextExt as _;
 use crate::math::{IeeeExt, apply_random_float_error_ulp};
 use crate::*;
 
+/// Check that the number of args is what we expect.
+fn check_intrinsic_arg_count<'a, 'tcx, const N: usize>(
+    args: &'a [OpTy<'tcx>],
+) -> InterpResult<'tcx, &'a [OpTy<'tcx>; N]>
+where
+    &'a [OpTy<'tcx>; N]: TryFrom<&'a [OpTy<'tcx>]>,
+{
+    if let Ok(ops) = args.try_into() {
+        return interp_ok(ops);
+    }
+    throw_ub_format!(
+        "incorrect number of arguments for intrinsic: got {}, expected {}",
+        args.len(),
+        N
+    )
+}
+
 impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
 pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     fn call_intrinsic(
@@ -114,7 +131,8 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
                 ));
             }
             "catch_unwind" => {
-                this.handle_catch_unwind(args, dest, ret)?;
+                let [try_fn, data, catch_fn] = check_intrinsic_arg_count(args)?;
+                this.handle_catch_unwind(try_fn, data, catch_fn, dest, ret)?;
                 // This pushed a stack frame, don't jump to `ret`.
                 return interp_ok(EmulateItemResult::AlreadyJumped);
             }

src/tools/miri/src/intrinsics/simd.rs

@@ -6,9 +6,8 @@ use rustc_middle::ty::FloatTy;
 use rustc_middle::{mir, ty};
 use rustc_span::{Symbol, sym};
 
-use crate::helpers::{
-    ToHost, ToSoft, bool_to_simd_element, check_intrinsic_arg_count, simd_element_to_bool,
-};
+use super::check_intrinsic_arg_count;
+use crate::helpers::{ToHost, ToSoft, bool_to_simd_element, simd_element_to_bool};
 use crate::*;
 
 #[derive(Copy, Clone)]

src/tools/miri/src/lib.rs

@@ -7,6 +7,7 @@
 #![feature(never_type)]
 #![feature(try_blocks)]
 #![feature(io_error_more)]
+#![feature(if_let_guard)]
 #![feature(variant_count)]
 #![feature(yeet_expr)]
 #![feature(nonzero_ops)]
@@ -158,6 +159,7 @@ pub use crate::shims::foreign_items::{DynSym, EvalContextExt as _};
 pub use crate::shims::io_error::{EvalContextExt as _, IoError, LibcError};
 pub use crate::shims::os_str::EvalContextExt as _;
 pub use crate::shims::panic::EvalContextExt as _;
+pub use crate::shims::sig::EvalContextExt as _;
 pub use crate::shims::time::EvalContextExt as _;
 pub use crate::shims::tls::TlsData;
 pub use crate::shims::unwind::{CatchUnwindData, EvalContextExt as _};

src/tools/miri/src/shims/aarch64.rs

@@ -20,7 +20,7 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let unprefixed_name = link_name.as_str().strip_prefix("llvm.aarch64.").unwrap();
         match unprefixed_name {
             "isb" => {
-                let [arg] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
+                let [arg] = this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
                 let arg = this.read_scalar(arg)?.to_i32()?;
                 match arg {
                     // SY ("full system scope")
@@ -38,7 +38,8 @@ pub(super) trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             // `left` input, the second half of the output from the `right` input.
             // https://developer.arm.com/architectures/instruction-sets/intrinsics/vpmaxq_u8
             "neon.umaxp.v16i8" => {
-                let [left, right] = this.check_shim(abi, CanonAbi::C, link_name, args)?;
+                let [left, right] =
+                    this.check_shim_sig_lenient(abi, CanonAbi::C, link_name, args)?;
 
                 let (left, left_len) = this.project_to_simd(left)?;
                 let (right, right_len) = this.project_to_simd(right)?;