Fix closure untupling/retupling (#373)

Use the rust-call ABI and MIR-level spread args flag to untuple/retuple closures and shims.

Co-authored-by: Adrian Palacios <73246657+adpaco-aws@users.noreply.github.com>

compiler/rustc_codegen_llvm/src/gotoc/metadata.rs

@@ -255,13 +255,37 @@ impl<'tcx> GotocCtx<'tcx> {
         self.monomorphize(p.ty(self.current_fn().mir().local_decls(), self.tcx).ty)
     }
 
-    pub fn closure_params(&self, substs: ty::subst::SubstsRef<'tcx>) -> Vec<Ty<'tcx>> {
-        let sig = self.monomorphize(substs.as_closure().sig());
-        let args = match sig.skip_binder().inputs()[0].kind() {
-            ty::Tuple(substs) => substs.iter().map(|s| s.expect_ty()),
-            _ => unreachable!("this argument of a closure must be a tuple"),
-        };
-        args.collect()
+    /// Closures expect their last arg untupled at call site, see comment at
+    /// ty_needs_closure_untupled.
+    fn sig_with_closure_untupled(&self, sig: ty::PolyFnSig<'tcx>) -> ty::PolyFnSig<'tcx> {
+        debug!("sig_with_closure_untupled sig: {:?}", sig);
+        let fn_sig = sig.skip_binder();
+        if let Some((tupe, prev_args)) = fn_sig.inputs().split_last() {
+            let args: Vec<Ty<'tcx>> = match tupe.kind() {
+                ty::Tuple(substs) => substs.iter().map(|s| s.expect_ty()),
+                _ => unreachable!("the final argument of a closure must be a tuple"),
+            }
+            .collect();
+
+            // The leading argument should be exactly the environment
+            assert!(prev_args.len() == 1);
+            let env = prev_args[0].clone();
+
+            // Recombine arguments: environment first, then the flattened tuple elements
+            let recombined_args = iter::once(env).chain(args);
+
+            return ty::Binder::bind_with_vars(
+                self.tcx.mk_fn_sig(
+                    recombined_args,
+                    fn_sig.output(),
+                    fn_sig.c_variadic,
+                    fn_sig.unsafety,
+                    fn_sig.abi,
+                ),
+                sig.bound_vars(),
+            );
+        }
+        sig
     }
 
     fn closure_sig(
@@ -283,30 +307,41 @@ impl<'tcx> GotocCtx<'tcx> {
         let env_region = ty::ReLateBound(ty::INNERMOST, br);
         let env_ty = self.tcx.closure_env_ty(def_id, substs, env_region).unwrap();
 
-        // The parameter types are tupled, but we want to have them in a vector
-        let params = self.closure_params(substs);
         let sig = sig.skip_binder();
 
         // We build a binder from `sig` where:
         //  * `inputs` contains a sequence with the closure and parameter types
         //  * the rest of attributes are obtained from `sig`
-        ty::Binder::bind_with_vars(
+        let sig = ty::Binder::bind_with_vars(
             self.tcx.mk_fn_sig(
-                iter::once(env_ty).chain(params.iter().cloned()),
+                iter::once(env_ty).chain(iter::once(sig.inputs()[0])),
                 sig.output(),
                 sig.c_variadic,
                 sig.unsafety,
                 sig.abi,
             ),
             bound_vars,
-        )
+        );
+
+        // The parameter types are tupled, but we want to have them in a vector
+        self.sig_with_closure_untupled(sig)
     }
 
     pub fn fn_sig_of_instance(&self, instance: Instance<'tcx>) -> ty::PolyFnSig<'tcx> {
         let fntyp = instance.ty(self.tcx, ty::ParamEnv::reveal_all());
         self.monomorphize(match fntyp.kind() {
             ty::Closure(def_id, subst) => self.closure_sig(*def_id, subst),
-            _ => fntyp.fn_sig(self.tcx),
+            ty::FnPtr(..) | ty::FnDef(..) => {
+                let sig = fntyp.fn_sig(self.tcx);
+                // Some virtual calls through a vtable may actually be closures
+                // or shims that also need the arguments untupled, even though
+                // the kind of the trait type is not a ty::Closure.
+                if self.ty_needs_closure_untupled(fntyp) {
+                    return self.sig_with_closure_untupled(sig);
+                }
+                sig
+            }
+            _ => unreachable!("Can't get function signature of type: {:?}", fntyp),
         })
     }
 

compiler/rustc_codegen_llvm/src/gotoc/mod.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 use bitflags::_core::any::Any;
 use cbmc::goto_program::symtab_transformer;
-use cbmc::goto_program::{Stmt, Symbol, SymbolTable};
+use cbmc::goto_program::{Expr, Stmt, Symbol, SymbolTable};
 use cbmc::{MachineModel, RoundingMode};
 use metadata::*;
 use rustc_codegen_ssa::traits::CodegenBackend;
@@ -12,10 +12,10 @@ use rustc_hir::def_id::DefId;
 use rustc_middle::dep_graph::{WorkProduct, WorkProductId};
 use rustc_middle::middle::cstore::{EncodedMetadata, MetadataLoaderDyn};
 use rustc_middle::mir::mono::{CodegenUnit, MonoItem};
-use rustc_middle::mir::{BasicBlock, BasicBlockData, Body, HasLocalDecls};
+use rustc_middle::mir::{BasicBlock, BasicBlockData, Body, HasLocalDecls, Local};
 use rustc_middle::ty::print::with_no_trimmed_paths;
 use rustc_middle::ty::query::Providers;
-use rustc_middle::ty::{self, Instance, TyCtxt};
+use rustc_middle::ty::{self, Instance, TyCtxt, TyS};
 use rustc_serialize::json::ToJson;
 use rustc_session::config::{OutputFilenames, OutputType};
 use rustc_session::Session;
@@ -135,10 +135,12 @@ impl<'tcx> GotocCtx<'tcx> {
         let mir = self.current_fn().mir();
         let ldecls = mir.local_decls();
         ldecls.indices().for_each(|lc| {
-            let base_name = match self.find_debug_info(&lc) {
-                None => format!("var_{}", lc.index()),
-                Some(info) => format!("{}", info.name),
-            };
+            if Some(lc) == mir.spread_arg {
+                // We have already added this local in the function prelude, so
+                // skip adding it again here.
+                return;
+            }
+            let base_name = self.codegen_var_base_name(&lc);
             let name = self.codegen_var_name(&lc);
             let ldata = &ldecls[lc];
             let t = self.monomorphize(ldata.ty);
@@ -149,13 +151,96 @@ impl<'tcx> GotocCtx<'tcx> {
             let sym_e = sym.to_expr();
             self.symbol_table.insert(sym);
 
-            // declare local variables. 0 represents return value
+            // Index 0 represents the return value, which does not need to be
+            // declared in the first block
             if lc.index() < 1 || lc.index() > mir.arg_count {
                 self.current_fn_mut().push_onto_block(Stmt::decl(sym_e, None, loc));
             }
         });
     }
 
+    /// MIR functions have a `spread_arg` field that specifies whether the
+    /// final argument to the function is "spread" at the LLVM/codegen level
+    /// from a tuple into its individual components. (Used for the "rust-
+    /// call" ABI, necessary because dynamic trait closure cannot have an
+    /// argument list in MIR that is both generic and variadic, so Rust
+    /// allows a generic tuple).
+    ///
+    /// If `spread_arg` is Some, then the wrapped value is the local that is
+    /// to be "spread"/untupled. However, the MIR function body itself expects
+    /// the tuple instead of the individual components, so we need to generate
+    /// a function prelude that _retuples_, that is, writes the components
+    /// back to the tuple local for use in the body.
+    ///
+    /// See:
+    /// https://rust-lang.zulipchat.com/#narrow/stream/182449-t-compiler.2Fhelp/topic/Determine.20untupled.20closure.20args.20from.20Instance.3F
+    fn codegen_function_prelude(&mut self) {
+        let mir = self.current_fn().mir();
+        if mir.spread_arg.is_none() {
+            // No special tuple argument, no work to be done.
+            return;
+        }
+        let spread_arg = mir.spread_arg.unwrap();
+        let spread_data = &mir.local_decls()[spread_arg];
+        let loc = self.codegen_span2(&spread_data.source_info.span);
+
+        // When we codegen the function signature elsewhere, we will codegen the
+        // untupled version. So, the tuple argument itself needs to have a
+        // symbol declared for it outside of the function signature, we do that
+        // here.
+        let tup_typ = self.codegen_ty(self.monomorphize(spread_data.ty));
+        let tup_sym = Symbol::variable(
+            self.codegen_var_name(&spread_arg),
+            self.codegen_var_base_name(&spread_arg),
+            tup_typ.clone(),
+            loc.clone(),
+        );
+        self.symbol_table.insert(tup_sym.clone());
+
+        // Get the function signature from MIR, _before_ we untuple
+        let fntyp = self.current_fn().instance().ty(self.tcx, ty::ParamEnv::reveal_all());
+        let sig = match fntyp.kind() {
+            ty::FnPtr(..) | ty::FnDef(..) => fntyp.fn_sig(self.tcx).skip_binder(),
+            // Closures themselves will have their arguments already untupled,
+            // see Zulip link above.
+            ty::Closure(..) => unreachable!(
+                "Unexpected `spread arg` set for closure, got: {:?}, {:?}",
+                fntyp,
+                self.current_fn().readable_name()
+            ),
+            _ => unreachable!(
+                "Expected function type for `spread arg` prelude, got: {:?}, {:?}",
+                fntyp,
+                self.current_fn().readable_name()
+            ),
+        };
+
+        // Now that we have the tuple, write the individual component locals
+        // back to it as a GotoC struct.
+        let tupe = sig.inputs().last().unwrap();
+        let args: Vec<&TyS<'tcx>> = match tupe.kind() {
+            ty::Tuple(substs) => substs.iter().map(|s| s.expect_ty()).collect(),
+            _ => unreachable!("a function's spread argument must be a tuple"),
+        };
+
+        // Convert each arg to a GotoC expression.
+        let mut arg_exprs = Vec::new();
+        let starting_idx = sig.inputs().len();
+        for (arg_i, arg_t) in args.iter().enumerate() {
+            // The components come at the end, so offset by the untupled length.
+            let lc = Local::from_usize(arg_i + starting_idx);
+            let (name, base_name) = self.codegen_spread_arg_name(&lc);
+            let sym = Symbol::variable(name, base_name, self.codegen_ty(arg_t), loc.clone());
+            self.symbol_table.insert(sym.clone());
+            arg_exprs.push(sym.to_expr());
+        }
+
+        // Finally, combine the expression into a struct.
+        let tuple_expr = Expr::struct_expr_from_values(tup_typ, arg_exprs, &self.symbol_table)
+            .with_location(loc.clone());
+        self.current_fn_mut().push_onto_block(Stmt::decl(tup_sym.to_expr(), Some(tuple_expr), loc));
+    }
+
     /// collect all labels for goto
     fn codegen_prepare_blocks(&self) -> Vec<String> {
         self.current_fn().mir().basic_blocks().indices().map(|bb| format!("{:?}", bb)).collect()
@@ -202,6 +287,7 @@ impl<'tcx> GotocCtx<'tcx> {
             self.print_instance(instance, mir);
             let labels = self.codegen_prepare_blocks();
             self.current_fn_mut().set_labels(labels);
+            self.codegen_function_prelude();
             self.codegen_declare_variables();
 
             mir.basic_blocks().iter_enumerated().for_each(|(bb, bbd)| self.codegen_block(bb, bbd));

compiler/rustc_codegen_llvm/src/gotoc/statement.rs

@@ -195,31 +195,50 @@ impl<'tcx> GotocCtx<'tcx> {
         }
     }
 
-    fn codegen_adjust_arguments(&mut self, fargs: &mut Vec<Expr>, instance: Instance<'tcx>) {
-        let ins_ty = self.monomorphize(instance.ty(self.tcx, ty::ParamEnv::reveal_all()));
-        debug!("codegen_adjust_arguments instance: {:?} ins_ty: {:?}", instance, ins_ty);
-        if let ty::Closure(_, substs) = ins_ty.kind() {
-            // a closure takes two argument:
-            // 0. a struct representing the environment
-            // 1. a tuple containing the parameters
-            //
-            // however, for some reason, Rust decides to generate a function which still
-            // takes the first argument as the environment struct, but the tuple of parameters
-            // are flattened as subsequent parameters.
-            // therefore, we have to project out the corresponding fields when we detect
-            // an invocation of a closure.
-            //
-            // In the case where the closure takes no paramaters, rustc appears to elide the
-            // second (empty) tuple. So we should only expand the tuple if its there.
-            if fargs.len() > 1 {
-                assert_eq!(fargs.len(), 2);
-                let tupe = fargs.remove(1);
-                for (i, t) in self.closure_params(substs).iter().enumerate() {
-                    if !t.is_unit() {
-                        // Access the tupled parameters through the `member` operation
-                        let index_param = tupe.clone().member(&i.to_string(), &self.symbol_table);
-                        fargs.push(index_param);
+    fn codegen_untuple_closure_args(
+        &mut self,
+        instance: Instance<'tcx>,
+        fargs: &mut Vec<Expr>,
+        last_mir_arg: Option<&Operand<'tcx>>,
+    ) {
+        debug!(
+            "codegen_untuple_closure_args instance: {:?}, fargs {:?}",
+            self.readable_instance_name(instance),
+            fargs
+        );
+        // A closure takes two arguments:
+        //     0. a struct representing the environment
+        //     1. a tuple containing the parameters
+        //
+        // However, for some reason, Rust decides to generate a function which still
+        // takes the first argument as the environment struct, but the tuple of parameters
+        // are flattened as subsequent parameters.
+        // Therefore, we have to project out the corresponding fields when we detect
+        // an invocation of a closure.
+        //
+        // Note: In some cases, the enviroment struct has type FnDef, so we skip it in
+        // ignore_var_ty. So the tuple is always the last arg, but it might be in the
+        // first or the second position.
+        if fargs.len() > 0 {
+            let tupe = fargs.remove(fargs.len() - 1);
+            let tupled_args: Vec<Type> = match self.operand_ty(last_mir_arg.unwrap()).kind() {
+                ty::Tuple(tupled_args) => {
+                    // The tuple needs to be added back for type checking even if empty
+                    if tupled_args.is_empty() {
+                        fargs.push(tupe);
+                        return;
                     }
+                    tupled_args.iter().map(|s| self.codegen_ty(s.expect_ty())).collect()
+                }
+                _ => unreachable!("Argument to function with Abi::RustCall is not a tuple"),
+            };
+
+            // Unwrap as needed
+            for (i, t) in tupled_args.iter().enumerate() {
+                if !t.is_unit() {
+                    // Access the tupled parameters through the `member` operation
+                    let index_param = tupe.clone().member(&i.to_string(), &self.symbol_table);
+                    fargs.push(index_param);
                 }
             }
         }
@@ -253,7 +272,10 @@ impl<'tcx> GotocCtx<'tcx> {
                     Instance::resolve(self.tcx, ty::ParamEnv::reveal_all(), *defid, subst)
                         .unwrap()
                         .unwrap();
-                self.codegen_adjust_arguments(&mut fargs, instance);
+
+                if self.ty_needs_closure_untupled(funct) {
+                    self.codegen_untuple_closure_args(instance, &mut fargs, args.last());
+                }
 
                 if let Some(hk) = self.hooks.hook_applies(self.tcx, instance) {
                     return hk.handle(