From 18d731ee59e9f1d4868b5bbb812b8075e073e527 Mon Sep 17 00:00:00 2001
From: Matthew Maurer <mmaurer@google.com>
Date: Mon, 19 Feb 2024 17:40:38 +0000
Subject: [PATCH] CFI: Generate super vtables explicitly

CFI shimming means they're not gauranteed to be pre-generated.
Traditionally, the base vtable has all the elements of the supertrait
vtable, and so visiting the base vtable implies you don't need to visit
the supertrait vtable. However, with CFI the base vtable entries will
have invocation type `dyn Child`, and the parent vtable will have
invocation type `dyn Parent`, so they aren't actually the same instance,
and both must be visited.
---
 compiler/rustc_monomorphize/src/collector.rs | 55 +++++++++++++++-----
 tests/ui/sanitizer/cfi-super-vtable.rs       | 39 ++++++++++++++
 2 files changed, 81 insertions(+), 13 deletions(-)
 create mode 100644 tests/ui/sanitizer/cfi-super-vtable.rs

diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index c888aaeca1467..f60bf062ce322 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -191,6 +191,7 @@ use rustc_span::source_map::{dummy_spanned, respan, Spanned};
 use rustc_span::symbol::{sym, Ident};
 use rustc_span::{Span, DUMMY_SP};
 use rustc_target::abi::Size;
+use std::iter;
 use std::path::PathBuf;
 
 use crate::errors::{
@@ -1201,23 +1202,51 @@ fn create_mono_items_for_vtable_methods<'tcx>(
             assert!(!poly_trait_ref.has_escaping_bound_vars());
 
             // Walk all methods of the trait, including those of its supertraits
-            let entries = tcx.vtable_entries(poly_trait_ref);
-            let methods = entries
-                .iter()
-                .filter_map(|entry| match entry {
+            for entry in tcx.vtable_entries(poly_trait_ref) {
+                match entry {
                     VtblEntry::MetadataDropInPlace
                     | VtblEntry::MetadataSize
                     | VtblEntry::MetadataAlign
-                    | VtblEntry::Vacant => None,
-                    VtblEntry::TraitVPtr(_) => {
-                        // all super trait items already covered, so skip them.
-                        None
+                    | VtblEntry::Vacant => (),
+                    VtblEntry::TraitVPtr(super_trait) => {
+                        // If CFI shims are enabled, the super_trait will use a different invocation
+                        // type than the instances selected for this trait. This means we must walk
+                        // the super_trait pointer visit its instances as well.
+                        if tcx.sess.cfi_shims() {
+                            let pep: ty::PolyExistentialPredicate<'tcx> =
+                                super_trait.map_bound(|t| {
+                                    ty::ExistentialPredicate::Trait(
+                                        ty::ExistentialTraitRef::erase_self_ty(tcx, t),
+                                    )
+                                });
+                            let existential_predicates = tcx
+                                .mk_poly_existential_predicates_from_iter(
+                                    iter::once(pep).chain(trait_ty.iter().skip(1)),
+                                );
+                            let super_trait_ty = Ty::new_dynamic(
+                                tcx,
+                                existential_predicates,
+                                tcx.lifetimes.re_erased,
+                                ty::Dyn,
+                            );
+
+                            create_mono_items_for_vtable_methods(
+                                tcx,
+                                super_trait_ty,
+                                impl_ty,
+                                source,
+                                output,
+                            );
+                        }
+                    }
+                    VtblEntry::Method(instance) => {
+                        let instance = instance.cfi_shim(tcx, invoke_trait);
+                        if should_codegen_locally(tcx, &instance) {
+                            output.push(create_fn_mono_item(tcx, instance, source));
+                        }
                     }
-                    VtblEntry::Method(instance) => Some(instance.cfi_shim(tcx, invoke_trait))
-                        .filter(|instance| should_codegen_locally(tcx, instance)),
-                })
-                .map(|item| create_fn_mono_item(tcx, item, source));
-            output.extend(methods);
+                }
+            }
         };
 
         // Also add the destructor.
diff --git a/tests/ui/sanitizer/cfi-super-vtable.rs b/tests/ui/sanitizer/cfi-super-vtable.rs
new file mode 100644
index 0000000000000..2764ffc5426b4
--- /dev/null
+++ b/tests/ui/sanitizer/cfi-super-vtable.rs
@@ -0,0 +1,39 @@
+// Check that super-traits with vptrs have their shims generated
+
+//@ needs-sanitizer-cfi
+//@ compile-flags: --crate-type=bin -Cprefer-dynamic=off -Clto -Zsanitizer=cfi
+//@ compile-flags: -C codegen-units=1 -C opt-level=0
+//@ build-pass
+
+trait Parent1 {
+    fn p1(&self);
+}
+
+trait Parent2 {
+    fn p2(&self);
+}
+
+// We need two parent traits to force the vtable upcasting code to choose to add a pointer to
+// another vtable to the child. This vtable is generated even if trait upcasting is not in use.
+trait Child : Parent1 + Parent2 {
+    fn c(&self);
+}
+
+struct Foo;
+
+impl Parent1 for Foo {
+    fn p1(&self) {}
+}
+
+impl Parent2 for Foo {
+    fn p2(&self) {}
+}
+
+impl Child for Foo {
+    fn c(&self) {}
+}
+
+fn main() {
+    let x = &Foo as &dyn Child;
+    x.c();
+}