Add support to casting closure to function pointer (rust-lang#2124)

Change function declaration and call to ignore arguments that are ZST, and implement the closure to function pointer by reifying it to its FnOnce implementation.

The casting of closures to function pointers rely on the fact that the closures do not capture anything. Those closures are represented by empty structures, which should get pruned from its argument list of the FnOnce::call_once implementation.

Author: celinval

docs/src/rust-feature-support.md

@@ -154,9 +154,7 @@ not formally defined which makes it harder to ensure that we can properly model
 all their use cases.
 
 In particular, there are some outstanding issues to note here:
- * Unimplemented `PointerCast::ClosureFnPointer` in
-   [#274](https://github.com/model-checking/kani/issues/274) and `Variant` case
-   in projections type in
+ * Sanity check `Variant` type in projections
    [#448](https://github.com/model-checking/kani/issues/448).
  * Unexpected fat pointer results in
    [#277](https://github.com/model-checking/kani/issues/277),

kani-compiler/src/codegen_cprover_gotoc/codegen/function.rs

@@ -12,7 +12,7 @@ use kani_queries::UserInput;
 use rustc_ast::Attribute;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::DefId;
-use rustc_middle::mir::{HasLocalDecls, Local};
+use rustc_middle::mir::{Body, HasLocalDecls, Local};
 use rustc_middle::ty::layout::FnAbiOf;
 use rustc_middle::ty::{self, Instance};
 use std::collections::BTreeMap;
@@ -47,14 +47,19 @@ impl<'tcx> GotocCtx<'tcx> {
             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);
-            let t = self.codegen_ty(t);
+            let var_ty = self.monomorphize(ldata.ty);
+            let var_type = self.codegen_ty(var_ty);
             let loc = self.codegen_span(&ldata.source_info.span);
             // Indices [1, N] represent the function parameters where N is the number of parameters.
-            let sym =
-                Symbol::variable(name, base_name, t, self.codegen_span(&ldata.source_info.span))
-                    .with_is_hidden(!ldata.is_user_variable())
-                    .with_is_parameter(idx > 0 && idx <= num_args);
+            // Except that ZST fields are not included as parameters.
+            let sym = Symbol::variable(
+                name,
+                base_name,
+                var_type,
+                self.codegen_span(&ldata.source_info.span),
+            )
+            .with_is_hidden(!ldata.is_user_variable())
+            .with_is_parameter((idx > 0 && idx <= num_args) && !self.is_zst(var_ty));
             let sym_e = sym.to_expr();
             self.symbol_table.insert(sym);
 
@@ -96,29 +101,41 @@ impl<'tcx> GotocCtx<'tcx> {
         self.reset_current_fn();
     }
 
+    /// Codegen changes required due to the function ABI.
+    /// We currently untuple arguments for RustCall ABI where the `spread_arg` is set.
+    fn codegen_function_prelude(&mut self) {
+        let mir = self.current_fn().mir();
+        if let Some(spread_arg) = mir.spread_arg {
+            self.codegen_spread_arg(mir, spread_arg);
+        }
+    }
+
     /// 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
+    /// call" ABI, necessary because the function traits and closures 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.
+    /// These tuples are used in the MIR to invoke a shim, and it's used in the shim body.
+    ///
+    /// The `spread_arg` represents the the local variable that is to be "spread"/untupled.
+    /// However, the function body itself may refer to the members of
+    /// the tuple instead of the individual spread parameters, so we need to add to the
+    /// function prelude code that _retuples_, that is, writes the arguments
+    /// back to a local tuple that can be used 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.
+    fn codegen_spread_arg(&mut self, mir: &Body<'tcx>, spread_arg: Local) {
+        tracing::debug!(current=?self.current_fn, "codegen_spread_arg");
+        let spread_data = &mir.local_decls()[spread_arg];
+        let tup_ty = self.monomorphize(spread_data.ty);
+        if self.is_zst(tup_ty) {
+            // No need to spread a ZST since it will be ignored.
             return;
         }
-        let spread_arg = mir.spread_arg.unwrap();
-        let spread_data = &mir.local_decls()[spread_arg];
+
         let loc = self.codegen_span(&spread_data.source_info.span);
 
         // Get the function signature from MIR, _before_ we untuple
@@ -159,7 +176,7 @@ impl<'tcx> GotocCtx<'tcx> {
         // };
         // ```
         // Note how the compiler has reordered the fields to improve packing.
-        let tup_typ = self.codegen_ty(self.monomorphize(spread_data.ty));
+        let tup_type = self.codegen_ty(tup_ty);
 
         // We need to marshall the arguments into the tuple
         // The arguments themselves have been tacked onto the explicit function paramaters by
@@ -192,7 +209,7 @@ impl<'tcx> GotocCtx<'tcx> {
                 let (name, base_name) = self.codegen_spread_arg_name(&lc);
                 let sym = Symbol::variable(name, base_name, self.codegen_ty(arg_t), loc)
                     .with_is_hidden(false)
-                    .with_is_parameter(true);
+                    .with_is_parameter(!self.is_zst(arg_t));
                 // The spread arguments are additional function paramaters that are patched in
                 // They are to the function signature added in the `fn_typ` function.
                 // But they were never added to the symbol table, which we currently do here.
@@ -204,12 +221,12 @@ impl<'tcx> GotocCtx<'tcx> {
                 (arg_i.to_string().intern(), sym.to_expr())
             }));
         let marshalled_tuple_value =
-            Expr::struct_expr(tup_typ.clone(), marshalled_tuple_fields, &self.symbol_table)
+            Expr::struct_expr(tup_type.clone(), marshalled_tuple_fields, &self.symbol_table)
                 .with_location(loc);
         self.declare_variable(
             self.codegen_var_name(&spread_arg),
             self.codegen_var_base_name(&spread_arg),
-            tup_typ,
+            tup_type,
             Some(marshalled_tuple_value),
             loc,
         );

kani-compiler/src/codegen_cprover_gotoc/codegen/intrinsic.rs

@@ -53,7 +53,7 @@ impl<'tcx> GotocCtx<'tcx> {
 
         if let Some(target) = target {
             let loc = self.codegen_span(&span);
-            let fargs = self.codegen_funcall_args(args);
+            let fargs = self.codegen_funcall_args(args, false);
             Stmt::block(
                 vec![
                     self.codegen_intrinsic(instance, fargs, destination, Some(span)),

kani-compiler/src/codegen_cprover_gotoc/codegen/rvalue.rs

@@ -640,13 +640,13 @@ impl<'tcx> GotocCtx<'tcx> {
     fn codegen_pointer_cast(
         &mut self,
         k: &PointerCast,
-        o: &Operand<'tcx>,
+        operand: &Operand<'tcx>,
         t: Ty<'tcx>,
         loc: Location,
     ) -> Expr {
-        debug!(cast=?k, op=?o, ?loc, "codegen_pointer_cast");
+        debug!(cast=?k, op=?operand, ?loc, "codegen_pointer_cast");
         match k {
-            PointerCast::ReifyFnPointer => match self.operand_ty(o).kind() {
+            PointerCast::ReifyFnPointer => match self.operand_ty(operand).kind() {
                 ty::FnDef(def_id, substs) => {
                     let instance =
                         Instance::resolve(self.tcx, ty::ParamEnv::reveal_all(), *def_id, substs)
@@ -658,17 +658,23 @@ impl<'tcx> GotocCtx<'tcx> {
                 }
                 _ => unreachable!(),
             },
-            PointerCast::UnsafeFnPointer => self.codegen_operand(o),
+            PointerCast::UnsafeFnPointer => self.codegen_operand(operand),
             PointerCast::ClosureFnPointer(_) => {
-                let dest_typ = self.codegen_ty(t);
-                self.codegen_unimplemented_expr(
-                    "PointerCast::ClosureFnPointer",
-                    dest_typ,
-                    loc,
-                    "https://github.com/model-checking/kani/issues/274",
-                )
+                if let ty::Closure(def_id, substs) = self.operand_ty(operand).kind() {
+                    let instance = Instance::resolve_closure(
+                        self.tcx,
+                        *def_id,
+                        substs,
+                        ty::ClosureKind::FnOnce,
+                    )
+                    .expect("failed to normalize and resolve closure during codegen")
+                    .polymorphize(self.tcx);
+                    self.codegen_func_expr(instance, None).address_of()
+                } else {
+                    unreachable!("{:?} cannot be cast to a fn ptr", operand)
+                }
             }
-            PointerCast::MutToConstPointer => self.codegen_operand(o),
+            PointerCast::MutToConstPointer => self.codegen_operand(operand),
             PointerCast::ArrayToPointer => {
                 // TODO: I am not sure whether it is correct or not.
                 //
@@ -677,11 +683,11 @@ impl<'tcx> GotocCtx<'tcx> {
                 // if we had to, then [o] necessarily has type [T; n] where *T is a fat pointer, meaning
                 // T is either [T] or str. but neither type is sized, which shouldn't participate in
                 // codegen.
-                match self.operand_ty(o).kind() {
+                match self.operand_ty(operand).kind() {
                     ty::RawPtr(ty::TypeAndMut { ty, .. }) => {
                         // ty must be an array
                         if let ty::Array(_, _) = ty.kind() {
-                            let oe = self.codegen_operand(o);
+                            let oe = self.codegen_operand(operand);
                             oe.dereference() // : struct [T; n]
                                 .member("0", &self.symbol_table) // : T[n]
                                 .array_to_ptr() // : T*
@@ -693,8 +699,8 @@ impl<'tcx> GotocCtx<'tcx> {
                 }
             }
             PointerCast::Unsize => {
-                let src_goto_expr = self.codegen_operand(o);
-                let src_mir_type = self.operand_ty(o);
+                let src_goto_expr = self.codegen_operand(operand);
+                let src_mir_type = self.operand_ty(operand);
                 let dst_mir_type = t;
                 self.codegen_unsized_cast(src_goto_expr, src_mir_type, dst_mir_type)
             }

kani-compiler/src/codegen_cprover_gotoc/codegen/statement.rs

@@ -33,9 +33,8 @@ impl<'tcx> GotocCtx<'tcx> {
             StatementKind::Assign(box (l, r)) => {
                 let lty = self.place_ty(l);
                 let rty = self.rvalue_ty(r);
-                let llayout = self.layout_of(lty);
                 // we ignore assignment for all zero size types
-                if llayout.is_zst() {
+                if self.is_zst(lty) {
                     Stmt::skip(location)
                 } else if lty.is_fn_ptr() && rty.is_fn() && !rty.is_fn_ptr() {
                     // implicit address of a function pointer, e.g.
@@ -402,10 +401,19 @@ impl<'tcx> GotocCtx<'tcx> {
         }
     }
 
-    /// As part of **calling** a function (closure actually), we may need to un-tuple arguments.
+    /// As part of **calling** a function (or closure), we may need to un-tuple arguments.
     ///
-    /// See [GotocCtx::ty_needs_closure_untupled]
-    fn codegen_untuple_closure_args(
+    /// This function will replace the last `fargs` argument by its un-tupled version.
+    ///
+    /// Some context: A closure / shim takes two arguments:
+    ///     0. a struct (or a pointer to) representing the environment
+    ///     1. a tuple containing the parameters (if not empty)
+    ///
+    /// However, Rust generates a function where the tuple of parameters are flattened
+    /// as subsequent parameters.
+    ///
+    /// See [GotocCtx::ty_needs_untupled_args] for more details.
+    fn codegen_untupled_args(
         &mut self,
         instance: Instance<'tcx>,
         fargs: &mut Vec<Expr>,
@@ -416,32 +424,22 @@ impl<'tcx> GotocCtx<'tcx> {
             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 environment 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.
-        // Note 2: For empty closures, the only argument needed is the environment struct.
         if !fargs.is_empty() {
+            let tuple_ty = self.operand_ty(last_mir_arg.unwrap());
+            if self.is_zst(tuple_ty) {
+                // Don't pass anything if all tuple elements are ZST.
+                // ZST arguments are ignored.
+                return;
+            }
             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) => tupled_args.iter().map(|s| self.codegen_ty(s)).collect(),
-                _ => unreachable!("Argument to function with Abi::RustCall is not a tuple"),
-            };
-
-            // Unwrap as needed
-            for (i, _) in tupled_args.iter().enumerate() {
-                // Access the tupled parameters through the `member` operation
-                let index_param = tupe.clone().member(&i.to_string(), &self.symbol_table);
-                fargs.push(index_param);
+            if let ty::Tuple(tupled_args) = tuple_ty.kind() {
+                for (idx, arg_ty) in tupled_args.iter().enumerate() {
+                    if !self.is_zst(arg_ty) {
+                        // Access the tupled parameters through the `member` operation
+                        let idx_expr = tupe.clone().member(&idx.to_string(), &self.symbol_table);
+                        fargs.push(idx_expr);
+                    }
+                }
             }
         }
     }
@@ -459,16 +457,30 @@ impl<'tcx> GotocCtx<'tcx> {
     /// Generate Goto-C for each argument to a function call.
     ///
     /// N.B. public only because instrinsics use this directly, too.
-    pub(crate) fn codegen_funcall_args(&mut self, args: &[Operand<'tcx>]) -> Vec<Expr> {
-        args.iter()
-            .map(|o| {
-                if self.operand_ty(o).is_bool() {
-                    self.codegen_operand(o).cast_to(Type::c_bool())
+    /// When `skip_zst` is set to `true`, the return value will not include any argument that is ZST.
+    /// This is used because we ignore ZST arguments, except for intrinsics.
+    pub(crate) fn codegen_funcall_args(
+        &mut self,
+        args: &[Operand<'tcx>],
+        skip_zst: bool,
+    ) -> Vec<Expr> {
+        let fargs = args
+            .iter()
+            .filter_map(|o| {
+                let op_ty = self.operand_ty(o);
+                if op_ty.is_bool() {
+                    Some(self.codegen_operand(o).cast_to(Type::c_bool()))
+                } else if !self.is_zst(op_ty) || !skip_zst {
+                    Some(self.codegen_operand(o))
                 } else {
-                    self.codegen_operand(o)
+                    // We ignore ZST types.
+                    debug!(arg=?o, "codegen_funcall_args ignore");
+                    None
                 }
             })
-            .collect()
+            .collect();
+        debug!(?fargs, "codegen_funcall_args");
+        fargs
     }
 
     /// Generates Goto-C for a MIR [TerminatorKind::Call] statement.
@@ -498,16 +510,17 @@ impl<'tcx> GotocCtx<'tcx> {
 
         let loc = self.codegen_span(&span);
         let funct = self.operand_ty(func);
-        let mut fargs = self.codegen_funcall_args(args);
+        let mut fargs = self.codegen_funcall_args(args, true);
         match &funct.kind() {
             ty::FnDef(defid, subst) => {
                 let instance =
                     Instance::resolve(self.tcx, ty::ParamEnv::reveal_all(), *defid, subst)
                         .unwrap()
                         .unwrap();
 
-                if self.ty_needs_closure_untupled(funct) {
-                    self.codegen_untuple_closure_args(instance, &mut fargs, args.last());
+                // TODO(celina): Move this check to be inside codegen_funcall_args.
+                if self.ty_needs_untupled_args(funct) {
+                    self.codegen_untupled_args(instance, &mut fargs, args.last());
                 }
 
                 if let Some(hk) = self.hooks.hook_applies(self.tcx, instance) {
@@ -608,6 +621,7 @@ impl<'tcx> GotocCtx<'tcx> {
     ) -> Vec<Stmt> {
         let vtable_field_name = self.vtable_field_name(def_id, idx);
         trace!(?self_ty, ?place, ?vtable_field_name, "codegen_virtual_funcall");
+        debug!(?fargs, "codegen_virtual_funcall");
 
         let trait_fat_ptr = self.extract_ptr(fargs[0].clone(), self_ty);
         assert!(