Add support to all object-safe receiver types (#1326)

* Add support to all object-safe receiver types

For object-safe traits, associated methods that can be dynamically
dispatched must have the receiver type of Self as one of the following:

- &Self
- &mut Self
- Box<Self>
- Rc<Self>
- Arc<Self>
- Pin<P> where P is one of the types above

ref: https://doc.rust-lang.org/reference/items/traits.html#object-safety

Before this change, we only supported &Self and &mut Self. In order to
add support to the other types, which are structures not pointers, we
needed to make sure that the structure type that we encode in the vtable
is compatible with the concrete type (ref 6.2.7.3 ISO/IEC 9899:202x).
For that, we mirror the original type and replace the raw pointer type
to be the data pointer type of `&dyn T`.

* Move test to expected

Issues with vtable encoding manifest as unreachable assertions instead
of verification failures. Thus, use expected test to ensure the
assertions are reachable and successful.

Author: celinval

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

@@ -21,7 +21,7 @@ use rustc_middle::ty::layout::LayoutOf;
 use rustc_middle::ty::{Instance, InstanceDef, Ty};
 use rustc_span::Span;
 use rustc_target::abi::{FieldsShape, Primitive, TagEncoding, Variants};
-use tracing::{debug, info_span, warn};
+use tracing::{debug, info_span, trace, warn};
 
 impl<'tcx> GotocCtx<'tcx> {
     fn codegen_ret_unit(&mut self) -> Stmt {
@@ -379,6 +379,7 @@ impl<'tcx> GotocCtx<'tcx> {
         target: &Option<BasicBlock>,
         span: Span,
     ) -> Stmt {
+        debug!(?func, ?args, ?destination, ?span, "codegen_funcall");
         if self.is_intrinsic(func) {
             return self.codegen_funcall_of_intrinsic(func, args, destination, target, span);
         }
@@ -411,19 +412,9 @@ impl<'tcx> GotocCtx<'tcx> {
                     }
                     // Handle a virtual function call via a vtable lookup
                     InstanceDef::Virtual(def_id, idx) => {
-                        // We must have at least one argument, and the first one
-                        // should be a fat pointer for the trait
-                        let trait_fat_ptr = fargs[0].to_owned();
-
-                        //Check the Gotoc-level fat pointer type
-                        assert!(
-                            trait_fat_ptr.typ().is_rust_trait_fat_ptr(&self.symbol_table),
-                            "Expected fat pointer, got:\n{:?}",
-                            trait_fat_ptr,
-                        );
-
+                        let self_ty = self.operand_ty(&args[0]);
                         self.codegen_virtual_funcall(
-                            trait_fat_ptr,
+                            self_ty,
                             def_id,
                             idx,
                             destination,
@@ -467,32 +458,87 @@ impl<'tcx> GotocCtx<'tcx> {
         };
     }
 
+    /// Extract a reference to self for virtual method calls.
+    /// See [codegen_dynamic_function_sig](GotocCtx::codegen_dynamic_function_sig) for more
+    /// details.
+    fn extract_ptr(&self, arg_expr: Expr, arg_ty: Ty<'tcx>) -> Expr {
+        // Generate an expression that indexes the pointer.
+        let expr = self
+            .receiver_data_path(arg_ty)
+            .fold(arg_expr, |curr_expr, (name, _)| curr_expr.member(name, &self.symbol_table));
+
+        trace!(?arg_ty, gotoc_ty=?expr.typ(), gotoc_expr=?expr.value(), "extract_ptr");
+        expr
+    }
+
+    /// Codegen the dynamic call to a trait method via the fat pointer vtable.
+    ///
+    /// If the original call was of the form
+    /// f(arg0, arg1);
+    ///
+    /// The new call should be of the form
+    /// arg0.vtable->f(arg0.data,arg1);
+    ///
+    /// For that, we do the following:
+    /// 1. Extract the fat pointer out of the first argument.
+    /// 2. Obtain the function pointer out of the fat pointer vtable.
+    /// 3. Change the first argument to only reference the data pointer (instead of the fat one).
+    ///     - When the receiver type is a `struct` we need to build a structure that mirrors
+    ///       the original one but uses a thin pointer instead.
+    /// 4. Generate the function call.
     fn codegen_virtual_funcall(
         &mut self,
-        trait_fat_ptr: Expr,
+        self_ty: Ty<'tcx>,
         def_id: DefId,
         idx: usize,
         place: &Place<'tcx>,
         fargs: &mut [Expr],
         loc: Location,
     ) -> Vec<Stmt> {
         let vtable_field_name = self.vtable_field_name(def_id, idx);
+        trace!(?self_ty, ?place, ?vtable_field_name, "codegen_virtual_funcall");
+
+        let trait_fat_ptr = self.extract_ptr(fargs[0].clone(), self_ty);
+        assert!(
+            trait_fat_ptr.typ().is_rust_trait_fat_ptr(&self.symbol_table),
+            "Expected fat pointer, but got {:?}",
+            trait_fat_ptr.typ()
+        );
 
-        // Now that we have all the stuff we need, we can actually build the dynamic call
-        // If the original call was of the form
-        // f(arg0, arg1);
-        // The new call should be of the form
-        // arg0.vtable->f(arg0.data,arg1);
         let vtable_ref = trait_fat_ptr.to_owned().member("vtable", &self.symbol_table);
         let vtable = vtable_ref.dereference();
         let fn_ptr = vtable.member(vtable_field_name, &self.symbol_table);
-
-        // Update the argument from arg0 to arg0.data
-        fargs[0] = trait_fat_ptr.to_owned().member("data", &self.symbol_table);
+        trace!(fn_typ=?fn_ptr.typ(), "codegen_virtual_funcall");
+
+        let data_ptr = trait_fat_ptr.to_owned().member("data", &self.symbol_table);
+        let mut ret_stmts = vec![];
+        fargs[0] = if self_ty.is_adt() {
+            // Generate a temp variable and assign its inner pointer to the fat_ptr.data.
+            match fn_ptr.typ() {
+                Type::Pointer { typ: box Type::Code { parameters, .. } } => {
+                    let param_typ = parameters.first().unwrap().typ();
+                    let tmp = self.gen_temp_variable(param_typ.clone(), loc).to_expr();
+                    debug!(?tmp,
+                        orig=?data_ptr.typ(),
+                        "codegen_virtual_funcall");
+                    ret_stmts.push(Stmt::decl(tmp.clone(), None, loc));
+                    ret_stmts.push(Stmt::assign(
+                        self.extract_ptr(tmp.clone(), self_ty),
+                        data_ptr,
+                        loc,
+                    ));
+                    tmp
+                }
+                _ => unreachable!("Unexpected virtual function type: {:?}", fn_ptr.typ()),
+            }
+        } else {
+            // Update the argument from arg0 to arg0.data if arg0 is a fat pointer.
+            data_ptr
+        };
 
         // For soundness, add an assertion that the vtable function call is not null.
-        // Otherwise, CBMC might treat this as an assert(0) and later user-added assertions
-        // could be vacuously true.
+        // Otherwise, CBMC might treat this as an assume(0) and later user-added assertions
+        // could become unreachable.
         let call_is_nonnull = fn_ptr.clone().is_nonnull();
         let assert_msg = format!("Non-null virtual function call for {:?}", vtable_field_name);
         let assert_nonnull =
@@ -506,7 +552,9 @@ impl<'tcx> GotocCtx<'tcx> {
         } else {
             call_stmt
         };
-        vec![assert_nonnull, call_stmt]
+        ret_stmts.push(assert_nonnull);
+        ret_stmts.push(call_stmt);
+        ret_stmts
     }
 
     /// A place is similar to the C idea of a LHS. For example, the returned value of a function call is stored to a place.

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

@@ -101,6 +101,7 @@ impl TypeExt for Type {
         }
     }
 }
+
 trait ExprExt {
     fn unit(symbol_table: &SymbolTable) -> Self;
 
@@ -354,9 +355,11 @@ impl<'tcx> GotocCtx<'tcx> {
         Type::unit().to_typedef(inner_name)
     }
 
-    /// This will codegen the raw pointer to the trait data.
+    /// Codegen the pointer type for a concrete object that implements the trait object.
+    /// I.e.: A trait object is a fat pointer which contains a pointer to a concrete object
+    /// and a pointer to its vtable. This method returns a type for the first pointer.
     pub fn codegen_trait_data_pointer(&mut self, typ: ty::Ty<'tcx>) -> Type {
-        assert!(typ.is_trait());
+        assert!(self.use_vtable_fat_pointer(typ));
         self.codegen_ty(typ).to_pointer()
     }
 
@@ -440,6 +443,7 @@ impl<'tcx> GotocCtx<'tcx> {
 
                 vtable_base.append(&mut flds);
             }
+            debug!(ty=?t, ?vtable_base, "trait_vtable_field_types");
             vtable_base
         } else {
             unreachable!("Expected to get a dynamic object here");
@@ -918,24 +922,35 @@ impl<'tcx> GotocCtx<'tcx> {
 
     /// Generate code for a trait function declaration.
     ///
-    /// Dynamic function calls first parameter is the fat-pointer representing self.
-    /// For closures, the type of the first argument is dyn T not &dyn T.
-    pub fn codegen_dynamic_function_sig(&mut self, sig: PolyFnSig<'tcx>) -> Type {
+    /// Dynamic function calls first parameter is self which must be one of the following:
+    ///
+    /// As of Jul 2022:
+    /// P = &Self | &mut Self | Box<Self> | Rc<Self> | Arc<Self>
+    /// S = P | Pin<P>
+    ///
+    /// See <https://doc.rust-lang.org/reference/items/traits.html#object-safety> for more details.
+    fn codegen_dynamic_function_sig(&mut self, sig: PolyFnSig<'tcx>) -> Type {
         let sig = self.monomorphize(sig);
         let sig = self.tcx.normalize_erasing_late_bound_regions(ty::ParamEnv::reveal_all(), sig);
-
         let mut is_first = true;
         let params = sig
             .inputs()
             .iter()
             .filter_map(|arg_type| {
                 if is_first {
-                    // This should &dyn T or dyn T (for closures).
                     is_first = false;
-                    let first_ty = pointee_type(*arg_type).unwrap_or(*arg_type);
                     debug!(self_type=?arg_type, fn_signature=?sig, "codegen_dynamic_function_sig");
-                    assert!(first_ty.is_trait(), "Expected self type to be a trait");
-                    Some(self.codegen_trait_data_pointer(first_ty))
+                    if arg_type.is_ref() {
+                        // Convert fat pointer to thin pointer to data portion.
+                        let first_ty = pointee_type(*arg_type).unwrap();
+                        Some(self.codegen_trait_data_pointer(first_ty))
+                    } else if arg_type.is_trait() {
+                        // Convert dyn T to thin pointer.
+                        Some(self.codegen_trait_data_pointer(*arg_type))
+                    } else {
+                        // Codegen type with thin pointer (E.g.: Box<dyn T> -> Box<data_ptr>).
+                        Some(self.codegen_trait_receiver(*arg_type))
+                    }
                 } else if self.ignore_var_ty(*arg_type) {
                     debug!("Ignoring type {:?} in function signature", arg_type);
                     None
@@ -1367,6 +1382,52 @@ impl<'tcx> GotocCtx<'tcx> {
             _ => false,
         }
     }
+
+    /// Generate code for a valid object-safe trait receiver type.
+    ///
+    /// Note that all these types only contain the data pointer and ZST fields. Thus, we generate
+    /// the non-ZST branch manually. In some cases, this method is called from inside
+    /// `codegen_ty(arg_ty)` so we don't have information about the final type.
+    fn codegen_trait_receiver(&mut self, arg_ty: Ty<'tcx>) -> Type {
+        // Collect structs that need to be modified
+        // Collect the non-ZST fields until we find a fat pointer.
+        let mut data_path = vec![arg_ty];
+        data_path.extend(self.receiver_data_path(arg_ty).map(|(_, typ)| typ));
+
+        trace!(?arg_ty, ?data_path, "codegen_trait_receiver");
+        let orig_pointer_ty = data_path.pop().unwrap();
+        assert!(self.is_vtable_fat_pointer(orig_pointer_ty));
+
+        // Traverse type and replace pointer type.
+        let ptr_type = self.codegen_trait_data_pointer(pointee_type(orig_pointer_ty).unwrap());
+        data_path.iter().rev().fold(ptr_type, |last_type, curr| {
+            // Codegen the type replacing the non-zst field.
+            let new_name = self.ty_mangled_name(*curr).to_string() + "::WithDataPtr";
+            if let ty::Adt(adt_def, adt_substs) = curr.kind() {
+                let fields = &adt_def.variants().get(VariantIdx::from_u32(0)).unwrap().fields;
+                self.ensure_struct(new_name, NO_PRETTY_NAME, |ctx, s_name| {
+                    let fields_shape = ctx.layout_of(*curr).layout.fields();
+                    let components = fields_shape
+                        .index_by_increasing_offset()
+                        .map(|idx| {
+                            let name = fields[idx].name.to_string().intern();
+                            let field_ty = fields[idx].ty(ctx.tcx, adt_substs);
+                            let typ = if !ctx.layout_of(field_ty).is_zst() {
+                                last_type.clone()
+                            } else {
+                                ctx.codegen_ty(field_ty)
+                            };
+                            DatatypeComponent::Field { name, typ }
+                        })
+                        .collect();
+                    trace!(?data_path, ?curr, ?s_name, ?components, "codegen_trait_receiver");
+                    components
+                })
+            } else {
+                unreachable!("Expected structs only {:?}", curr);
+            }
+        })
+    }
 }
 
 /// Use maps instead of lists to manage mir struct components.
@@ -1412,6 +1473,70 @@ impl<'tcx> GotocCtx<'tcx> {
         }
         if typ.is_trait() { Some(typ) } else { None }
     }
+
+    /// This function provides an iterator that traverses the data path of a receiver type. I.e.:
+    /// the path that leads to the data pointer.
+    ///
+    /// E.g.: For `Rc<dyn T>` where the Rc definition is:
+    /// ```
+    /// pub struct Rc<T: ?Sized> {
+    ///    ptr: NonNull<RcBox<T>>,
+    ///    phantom: PhantomData<RcBox<T>>,
+    /// }
+    ///
+    /// pub struct NonNull<T: ?Sized> {
+    ///    pointer: *const T,
+    /// }
+    /// ```
+    ///
+    /// The behavior will be:
+    /// ```text
+    /// let it = self.receiver_data_path(rc_typ);
+    /// assert_eq!(it.next(), Some((String::from("ptr"), non_null_typ);
+    /// assert_eq!(it.next(), Some((String::from("pointer"), raw_ptr_typ);
+    /// assert_eq!(it.next(), None);
+    /// ```
+    ///
+    /// Pre-condition: The argument must be a valid receiver for dispatchable trait functions.
+    /// See <https://doc.rust-lang.org/reference/items/traits.html#object-safety> for more details.
+    pub fn receiver_data_path<'a>(
+        self: &'a Self,
+        typ: Ty<'tcx>,
+    ) -> impl Iterator<Item = (String, Ty<'tcx>)> + 'a {
+        struct ReceiverIter<'tcx, 'a> {
+            pub curr: Ty<'tcx>,
+            pub ctx: &'a GotocCtx<'tcx>,
+        }
+
+        impl<'tcx, 'a> Iterator for ReceiverIter<'tcx, 'a> {
+            type Item = (String, Ty<'tcx>);
+
+            fn next(&mut self) -> Option<Self::Item> {
+                if let ty::Adt(adt_def, adt_substs) = self.curr.kind() {
+                    assert_eq!(
+                        adt_def.variants().len(),
+                        1,
+                        "Expected a single-variant ADT. Found {:?}",
+                        self.curr
+                    );
+                    let ctx = self.ctx;
+                    let fields = &adt_def.variants().get(VariantIdx::from_u32(0)).unwrap().fields;
+                    let mut non_zsts = fields
+                        .iter()
+                        .filter(|field| !ctx.layout_of(field.ty(ctx.tcx, adt_substs)).is_zst())
+                        .map(|non_zst| (non_zst.name.to_string(), non_zst.ty(ctx.tcx, adt_substs)));
+                    let (name, next) = non_zsts.next().expect("Expected one non-zst field.");
+                    self.curr = next;
+                    assert!(non_zsts.next().is_none(), "Expected only one non-zst field.");
+                    Some((name, self.curr))
+                } else {
+                    None
+                }
+            }
+        }
+
+        ReceiverIter { ctx: self, curr: typ }
+    }
 }
 
 /// The mir type is a mir pointer type.