simplify callee signature computation and make ABI handling more consistent

Author: RalfJung

compiler/rustc_const_eval/src/interpret/call.rs

@@ -353,23 +353,18 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
     ) -> InterpResult<'tcx> {
         let _trace = enter_trace_span!(M, step::init_stack_frame, %instance, tracing_separate_thread = Empty);
 
-        // A mismatch between caller and callee c-variadicness or fixed_count will error below.
-        let (fixed_count, callee_c_variadic_args) = if caller_fn_abi.c_variadic {
-            let callee_fn_sig = self.tcx.fn_sig(instance.def_id()).skip_binder();
-            let fixed_count = callee_fn_sig.inputs().skip_binder().len();
-
-            // NOTE: use the caller's fixed count here to slice into the caller's arguments. Using
-            // the callee's fixed_count could go out of bounds for malformed input programs.
-            let extra_tys = caller_fn_abi.args[caller_fn_abi.fixed_count as usize..]
-                .iter()
-                .map(|arg_abi| arg_abi.layout.ty);
-
-            (fixed_count, self.tcx.mk_type_list_from_iter(extra_tys))
+        // The first order of business is to figure out the callee signature.
+        // However, that requires the list of variadic arguments.
+        // We use the *caller* information to determine where to split the list of arguments,
+        // and then later check that the callee indeed has the same number of fixed arguments.
+        let extra_tys = if caller_fn_abi.c_variadic {
+            let fixed_count = usize::try_from(caller_fn_abi.fixed_count).unwrap();
+            let extra_tys = args[fixed_count..].iter().map(|arg| arg.layout().ty);
+            self.tcx.mk_type_list_from_iter(extra_tys)
         } else {
-            (caller_fn_abi.fixed_count.try_into().unwrap(), ty::List::empty())
+            ty::List::empty()
         };
-
-        let callee_fn_abi = self.fn_abi_of_instance(instance, callee_c_variadic_args)?;
+        let callee_fn_abi = self.fn_abi_of_instance(instance, extra_tys)?;
 
         if caller_fn_abi.conv != callee_fn_abi.conv {
             throw_ub_custom!(
@@ -387,7 +382,6 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
                 callee_is_c_variadic: callee_fn_abi.c_variadic,
             });
         }
-
         if caller_fn_abi.c_variadic && caller_fn_abi.fixed_count != callee_fn_abi.fixed_count {
             throw_ub!(CVariadicFixedCountMismatch {
                 caller: caller_fn_abi.fixed_count,
@@ -466,32 +460,20 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             // this is a single iterator (that handles `spread_arg`), then
             // `pass_argument` would be the loop body. It takes care to
             // not advance `caller_iter` for ignored arguments.
-            let mut callee_args_abis = if caller_fn_abi.c_variadic {
-                // Only the fixed arguments are passed normally. All remaining arguments are
-                // put into the va_list. Crucially, there can't be a caller location among those
-                // remaining arguments, so the caller location handling below does not interfere
-                // with the variadic argument handling: C-variadic functions cannot be
-                // `extern "Rust"` and `#[track_caller]` can only be applied to `extern "Rust"`.
-                assert!(!instance.def.requires_caller_location(*self.tcx));
-                callee_fn_abi.args[..fixed_count].iter().enumerate()
-            } else {
-                // NOTE: this iterator includes the extra caller location argument that is passed
-                // when `#[track_caller]` is used. The `fixed_count` does not account for this
-                // argument. This attribute is only allowed on `extern "Rust"` functions, so the
-                // c-variadic case does not need to handle the extra argument.
-                callee_fn_abi.args.iter().enumerate()
-            };
-
-            let mut it = body.args_iter().peekable();
-            while let Some(local) = it.next() {
+            let mut callee_args_abis = callee_fn_abi.args.iter().enumerate();
+            // Determine whether there is a special VaList argument. This is always the
+            // last argument, and since arguments start at index 1 that's `arg_count`.
+            let va_list_arg =
+                callee_fn_abi.c_variadic.then(|| mir::Local::from_usize(body.arg_count));
+            for local in body.args_iter() {
                 // Construct the destination place for this argument. At this point all
                 // locals are still dead, so we cannot construct a `PlaceTy`.
                 let dest = mir::Place::from(local);
                 // `layout_of_local` does more than just the instantiation we need to get the
                 // type, but the result gets cached so this avoids calling the instantiation
                 // query *again* the next time this local is accessed.
                 let ty = self.layout_of_local(self.frame(), local, None)?.ty;
-                if caller_fn_abi.c_variadic && it.peek().is_none() {
+                if Some(local) == va_list_arg {
                     // This is the last callee-side argument of a variadic function.
                     // This argument is a VaList holding the remaining caller-side arguments.
                     self.storage_live(local)?;
@@ -501,7 +483,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
 
                     // Consume the remaining arguments by putting them into the variable argument
                     // list.
-                    let varargs = self.allocate_varargs(&mut caller_args)?;
+                    let varargs = self.allocate_varargs(&mut caller_args, &mut callee_args_abis)?;
                     // When the frame is dropped, these variable arguments are deallocated.
                     self.frame_mut().va_list = varargs.clone();
                     let key = self.va_list_ptr(varargs.into());
@@ -554,7 +536,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             if instance.def.requires_caller_location(*self.tcx) {
                 callee_args_abis.next().unwrap();
             }
-            // Now we should have no more caller args or callee arg ABIs
+            // Now we should have no more caller args or callee arg ABIs.
             assert!(
                 callee_args_abis.next().is_none(),
                 "mismatch between callee ABI and callee body arguments"

compiler/rustc_const_eval/src/interpret/stack.rs

@@ -638,23 +638,30 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
 impl<'a, 'tcx: 'a, M: Machine<'tcx>> InterpCx<'tcx, M> {
     /// Consume the arguments provided by the iterator and store them as a list
     /// of variadic arguments. Return a list of the places that hold those arguments.
-    pub(crate) fn allocate_varargs<I>(
+    pub(crate) fn allocate_varargs<I, J>(
         &mut self,
         caller_args: &mut I,
+        callee_abis: &mut J,
     ) -> InterpResult<'tcx, Vec<MPlaceTy<'tcx, M::Provenance>>>
     where
         I: Iterator<Item = (&'a FnArg<'tcx, M::Provenance>, &'a ArgAbi<'tcx, Ty<'tcx>>)>,
+        J: Iterator<Item = (usize, &'a ArgAbi<'tcx, Ty<'tcx>>)>,
     {
         // Consume the remaining arguments and store them in fresh allocations.
         let mut varargs = Vec::new();
-        for (fn_arg, abi) in caller_args {
+        for (fn_arg, caller_abi) in caller_args {
+            // The callee ABI is entirely computed based on which arguments the caller has
+            // provided so it should not be possible to get a mismatch here.
+            let (_idx, callee_abi) = callee_abis.next().unwrap();
+            assert!(self.check_argument_compat(caller_abi, callee_abi)?);
             // FIXME: do we have to worry about in-place argument passing?
             let op = self.copy_fn_arg(fn_arg);
-            let mplace = self.allocate(abi.layout, MemoryKind::Stack)?;
+            let mplace = self.allocate(op.layout, MemoryKind::Stack)?;
             self.copy_op(&op, &mplace)?;
 
             varargs.push(mplace);
         }
+        assert!(callee_abis.next().is_none());
 
         interp_ok(varargs)
     }