Auto merge of #113856 - WaffleLapkin:vtablin', r=oli-obk

Refactor vtable encoding and optimize it for the case of multiple marker traits

This PR does two things
- Refactor `prepare_vtable_segments` (this was motivated by the other change, `prepare_vtable_segments` was quite hard to understand and while trying to edit it I've refactored it)
  - Mostly remove `loop`s labeled `break`s/`continue`s whenever there is a simpler solution
  - Also use `?`
- Make vtable format a bit more efficient wrt to marker traits
  - See the tests for an example

Fixes https://github.com/rust-lang/rust/issues/113840
cc `@crlf0710`

----

Review wise it's probably best to review each commit individually, as then it's more clear why the refactoring is correct.

I can split the last two commits (which change behavior) into a separate PR if it makes reviewing easier

compiler/rustc_infer/src/traits/util.rs

@@ -25,6 +25,13 @@ impl<'tcx> PredicateSet<'tcx> {
         Self { tcx, set: Default::default() }
     }
 
+    /// Adds a predicate to the set.
+    ///
+    /// Returns whether the predicate was newly inserted. That is:
+    /// - If the set did not previously contain this predicate, `true` is returned.
+    /// - If the set already contained this predicate, `false` is returned,
+    ///   and the set is not modified: original predicate is not replaced,
+    ///   and the predicate passed as argument is dropped.
     pub fn insert(&mut self, pred: ty::Predicate<'tcx>) -> bool {
         // We have to be careful here because we want
         //

compiler/rustc_trait_selection/src/traits/vtable.rs

@@ -24,8 +24,18 @@ pub enum VtblSegment<'tcx> {
 pub fn prepare_vtable_segments<'tcx, T>(
     tcx: TyCtxt<'tcx>,
     trait_ref: ty::PolyTraitRef<'tcx>,
-    mut segment_visitor: impl FnMut(VtblSegment<'tcx>) -> ControlFlow<T>,
+    segment_visitor: impl FnMut(VtblSegment<'tcx>) -> ControlFlow<T>,
 ) -> Option<T> {
+    prepare_vtable_segments_inner(tcx, trait_ref, segment_visitor).break_value()
+}
+
+/// Helper for [`prepare_vtable_segments`] that returns `ControlFlow`,
+/// such that we can use `?` in the body.
+fn prepare_vtable_segments_inner<'tcx, T>(
+    tcx: TyCtxt<'tcx>,
+    trait_ref: ty::PolyTraitRef<'tcx>,
+    mut segment_visitor: impl FnMut(VtblSegment<'tcx>) -> ControlFlow<T>,
+) -> ControlFlow<T> {
     // The following constraints holds for the final arrangement.
     // 1. The whole virtual table of the first direct super trait is included as the
     //    the prefix. If this trait doesn't have any super traits, then this step
@@ -71,20 +81,18 @@ pub fn prepare_vtable_segments<'tcx, T>(
     //  N, N-vptr, O
 
     // emit dsa segment first.
-    if let ControlFlow::Break(v) = (segment_visitor)(VtblSegment::MetadataDSA) {
-        return Some(v);
-    }
+    segment_visitor(VtblSegment::MetadataDSA)?;
 
     let mut emit_vptr_on_new_entry = false;
     let mut visited = PredicateSet::new(tcx);
     let predicate = trait_ref.without_const().to_predicate(tcx);
     let mut stack: SmallVec<[(ty::PolyTraitRef<'tcx>, _, _); 5]> =
-        smallvec![(trait_ref, emit_vptr_on_new_entry, None)];
+        smallvec![(trait_ref, emit_vptr_on_new_entry, maybe_iter(None))];
     visited.insert(predicate);
 
     // the main traversal loop:
     // basically we want to cut the inheritance directed graph into a few non-overlapping slices of nodes
-    // that each node is emitted after all its descendents have been emitted.
+    // such that each node is emitted after all its descendants have been emitted.
     // so we convert the directed graph into a tree by skipping all previously visited nodes using a visited set.
     // this is done on the fly.
     // Each loop run emits a slice - it starts by find a "childless" unvisited node, backtracking upwards, and it
@@ -105,80 +113,81 @@ pub fn prepare_vtable_segments<'tcx, T>(
     // Loop run #1: Emitting the slice [D C] (in reverse order). No one has a next-sibling node.
     // Loop run #1: Stack after exiting out is []. Now the function exits.
 
-    loop {
+    'outer: loop {
         // dive deeper into the stack, recording the path
         'diving_in: loop {
-            if let Some((inner_most_trait_ref, _, _)) = stack.last() {
-                let inner_most_trait_ref = *inner_most_trait_ref;
-                let mut direct_super_traits_iter = tcx
-                    .super_predicates_of(inner_most_trait_ref.def_id())
-                    .predicates
-                    .into_iter()
-                    .filter_map(move |(pred, _)| {
-                        pred.subst_supertrait(tcx, &inner_most_trait_ref).as_trait_clause()
-                    });
+            let &(inner_most_trait_ref, _, _) = stack.last().unwrap();
 
-                'diving_in_skip_visited_traits: loop {
-                    if let Some(next_super_trait) = direct_super_traits_iter.next() {
-                        if visited.insert(next_super_trait.to_predicate(tcx)) {
-                            // We're throwing away potential constness of super traits here.
-                            // FIXME: handle ~const super traits
-                            let next_super_trait = next_super_trait.map_bound(|t| t.trait_ref);
-                            stack.push((
-                                next_super_trait,
-                                emit_vptr_on_new_entry,
-                                Some(direct_super_traits_iter),
-                            ));
-                            break 'diving_in_skip_visited_traits;
-                        } else {
-                            continue 'diving_in_skip_visited_traits;
-                        }
-                    } else {
-                        break 'diving_in;
-                    }
+            let mut direct_super_traits_iter = tcx
+                .super_predicates_of(inner_most_trait_ref.def_id())
+                .predicates
+                .into_iter()
+                .filter_map(move |(pred, _)| {
+                    pred.subst_supertrait(tcx, &inner_most_trait_ref).as_trait_clause()
+                });
+
+            // Find an unvisited supertrait
+            match direct_super_traits_iter
+                .find(|&super_trait| visited.insert(super_trait.to_predicate(tcx)))
+            {
+                // Push it to the stack for the next iteration of 'diving_in to pick up
+                Some(unvisited_super_trait) => {
+                    // We're throwing away potential constness of super traits here.
+                    // FIXME: handle ~const super traits
+                    let next_super_trait = unvisited_super_trait.map_bound(|t| t.trait_ref);
+                    stack.push((
+                        next_super_trait,
+                        emit_vptr_on_new_entry,
+                        maybe_iter(Some(direct_super_traits_iter)),
+                    ))
                 }
+
+                // There are no more unvisited direct super traits, dive-in finished
+                None => break 'diving_in,
             }
         }
 
-        // Other than the left-most path, vptr should be emitted for each trait.
-        emit_vptr_on_new_entry = true;
-
         // emit innermost item, move to next sibling and stop there if possible, otherwise jump to outer level.
-        'exiting_out: loop {
-            if let Some((inner_most_trait_ref, emit_vptr, siblings_opt)) = stack.last_mut() {
-                if let ControlFlow::Break(v) = (segment_visitor)(VtblSegment::TraitOwnEntries {
-                    trait_ref: *inner_most_trait_ref,
-                    emit_vptr: *emit_vptr,
-                }) {
-                    return Some(v);
-                }
+        while let Some((inner_most_trait_ref, emit_vptr, mut siblings)) = stack.pop() {
+            segment_visitor(VtblSegment::TraitOwnEntries {
+                trait_ref: inner_most_trait_ref,
+                emit_vptr,
+            })?;
 
-                'exiting_out_skip_visited_traits: loop {
-                    if let Some(siblings) = siblings_opt {
-                        if let Some(next_inner_most_trait_ref) = siblings.next() {
-                            if visited.insert(next_inner_most_trait_ref.to_predicate(tcx)) {
-                                // We're throwing away potential constness of super traits here.
-                                // FIXME: handle ~const super traits
-                                let next_inner_most_trait_ref =
-                                    next_inner_most_trait_ref.map_bound(|t| t.trait_ref);
-                                *inner_most_trait_ref = next_inner_most_trait_ref;
-                                *emit_vptr = emit_vptr_on_new_entry;
-                                break 'exiting_out;
-                            } else {
-                                continue 'exiting_out_skip_visited_traits;
-                            }
-                        }
-                    }
-                    stack.pop();
-                    continue 'exiting_out;
-                }
+            // If we've emitted (fed to `segment_visitor`) a trait that has methods present in the vtable,
+            // we'll need to emit vptrs from now on.
+            if !emit_vptr_on_new_entry
+                && has_own_existential_vtable_entries(tcx, inner_most_trait_ref.def_id())
+            {
+                emit_vptr_on_new_entry = true;
+            }
+
+            if let Some(next_inner_most_trait_ref) =
+                siblings.find(|&sibling| visited.insert(sibling.to_predicate(tcx)))
+            {
+                // We're throwing away potential constness of super traits here.
+                // FIXME: handle ~const super traits
+                let next_inner_most_trait_ref =
+                    next_inner_most_trait_ref.map_bound(|t| t.trait_ref);
+
+                stack.push((next_inner_most_trait_ref, emit_vptr_on_new_entry, siblings));
+
+                // just pushed a new trait onto the stack, so we need to go through its super traits
+                continue 'outer;
             }
-            // all done
-            return None;
         }
+
+        // the stack is empty, all done
+        return ControlFlow::Continue(());
     }
 }
 
+/// Turns option of iterator into an iterator (this is just flatten)
+fn maybe_iter<I: Iterator>(i: Option<I>) -> impl Iterator<Item = I::Item> {
+    // Flatten is bad perf-vise, we could probably implement a special case here that is better
+    i.into_iter().flatten()
+}
+
 fn dump_vtable_entries<'tcx>(
     tcx: TyCtxt<'tcx>,
     sp: Span,
@@ -192,11 +201,23 @@ fn dump_vtable_entries<'tcx>(
     });
 }
 
+fn has_own_existential_vtable_entries(tcx: TyCtxt<'_>, trait_def_id: DefId) -> bool {
+    own_existential_vtable_entries_iter(tcx, trait_def_id).next().is_some()
+}
+
 fn own_existential_vtable_entries(tcx: TyCtxt<'_>, trait_def_id: DefId) -> &[DefId] {
+    tcx.arena.alloc_from_iter(own_existential_vtable_entries_iter(tcx, trait_def_id))
+}
+
+fn own_existential_vtable_entries_iter(
+    tcx: TyCtxt<'_>,
+    trait_def_id: DefId,
+) -> impl Iterator<Item = DefId> + '_ {
     let trait_methods = tcx
         .associated_items(trait_def_id)
         .in_definition_order()
         .filter(|item| item.kind == ty::AssocKind::Fn);
+
     // Now list each method's DefId (for within its trait).
     let own_entries = trait_methods.filter_map(move |&trait_method| {
         debug!("own_existential_vtable_entry: trait_method={:?}", trait_method);
@@ -211,7 +232,7 @@ fn own_existential_vtable_entries(tcx: TyCtxt<'_>, trait_def_id: DefId) -> &[Def
         Some(def_id)
     });
 
-    tcx.arena.alloc_from_iter(own_entries.into_iter())
+    own_entries
 }
 
 /// Given a trait `trait_ref`, iterates the vtable entries

tests/ui/traits/object/print_vtable_sizes.rs

@@ -10,8 +10,9 @@ trait C {
     fn x() {} // not object safe, shouldn't be reported
 }
 
-// This ideally should not have any upcasting cost,
-// but currently does due to a bug
+// This does not have any upcasting cost,
+// because both `Send` and `Sync` are traits
+// with no methods
 trait D: Send + Sync + help::MarkerWithSuper {}
 
 // This can't have no cost without reordering,

tests/ui/traits/object/print_vtable_sizes.stdout

@@ -1,8 +1,8 @@
-print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "D", "entries": "7", "entries_ignoring_upcasting": "4", "entries_for_upcasting": "3", "upcasting_cost_percent": "75" }
 print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "E", "entries": "6", "entries_ignoring_upcasting": "4", "entries_for_upcasting": "2", "upcasting_cost_percent": "50" }
 print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "G", "entries": "14", "entries_ignoring_upcasting": "11", "entries_for_upcasting": "3", "upcasting_cost_percent": "27.27272727272727" }
 print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "A", "entries": "6", "entries_ignoring_upcasting": "5", "entries_for_upcasting": "1", "upcasting_cost_percent": "20" }
 print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "B", "entries": "4", "entries_ignoring_upcasting": "4", "entries_for_upcasting": "0", "upcasting_cost_percent": "0" }
+print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "D", "entries": "4", "entries_ignoring_upcasting": "4", "entries_for_upcasting": "0", "upcasting_cost_percent": "0" }
 print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "F", "entries": "6", "entries_ignoring_upcasting": "6", "entries_for_upcasting": "0", "upcasting_cost_percent": "0" }
 print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "_::S", "entries": "3", "entries_ignoring_upcasting": "3", "entries_for_upcasting": "0", "upcasting_cost_percent": "0" }
 print-vtable-sizes { "crate_name": "<UNKNOWN_CRATE>", "trait_name": "_::S", "entries": "3", "entries_ignoring_upcasting": "3", "entries_for_upcasting": "0", "upcasting_cost_percent": "0" }

tests/ui/traits/vtable/multiple-markers.rs

@@ -0,0 +1,47 @@
+// Regression test for <https://github.com/rust-lang/rust/issues/113840>
+//
+// This test makes sure that multiple marker (method-less) traits can reuse the
+// same pointer for upcasting.
+//
+// build-fail
+#![crate_type = "lib"]
+#![feature(rustc_attrs)]
+
+// Markers
+trait M0 {}
+trait M1 {}
+trait M2 {}
+
+// Just a trait with a method
+trait T {
+    fn method(&self) {}
+}
+
+#[rustc_dump_vtable]
+trait A: M0 + M1 + M2 + T {} //~ error: vtable entries for `<S as A>`:
+
+#[rustc_dump_vtable]
+trait B: M0 + M1 + T + M2 {} //~ error: vtable entries for `<S as B>`:
+
+#[rustc_dump_vtable]
+trait C: M0 + T + M1 + M2 {} //~ error: vtable entries for `<S as C>`:
+
+#[rustc_dump_vtable]
+trait D: T + M0 + M1 + M2 {} //~ error: vtable entries for `<S as D>`:
+
+struct S;
+
+impl M0 for S {}
+impl M1 for S {}
+impl M2 for S {}
+impl T for S {}
+impl A for S {}
+impl B for S {}
+impl C for S {}
+impl D for S {}
+
+pub fn require_vtables() {
+    fn require_vtables(_: &dyn A, _: &dyn B, _: &dyn C, _: &dyn D) {}
+
+    require_vtables(&S, &S, &S, &S)
+}

tests/ui/traits/vtable/multiple-markers.stderr

@@ -0,0 +1,52 @@
+error: vtable entries for `<S as A>`: [
+           MetadataDropInPlace,
+           MetadataSize,
+           MetadataAlign,
+           Method(<S as T>::method),
+       ]
+  --> $DIR/multiple-markers.rs:21:1
+   |
+LL | trait A: M0 + M1 + M2 + T {}
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: vtable entries for `<S as B>`: [
+           MetadataDropInPlace,
+           MetadataSize,
+           MetadataAlign,
+           Method(<S as T>::method),
+           TraitVPtr(<S as M2>),
+       ]
+  --> $DIR/multiple-markers.rs:24:1
+   |
+LL | trait B: M0 + M1 + T + M2 {}
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: vtable entries for `<S as C>`: [
+           MetadataDropInPlace,
+           MetadataSize,
+           MetadataAlign,
+           Method(<S as T>::method),
+           TraitVPtr(<S as M1>),
+           TraitVPtr(<S as M2>),
+       ]
+  --> $DIR/multiple-markers.rs:27:1
+   |
+LL | trait C: M0 + T + M1 + M2 {}
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: vtable entries for `<S as D>`: [
+           MetadataDropInPlace,
+           MetadataSize,
+           MetadataAlign,
+           Method(<S as T>::method),
+           TraitVPtr(<S as M0>),
+           TraitVPtr(<S as M1>),
+           TraitVPtr(<S as M2>),
+       ]
+  --> $DIR/multiple-markers.rs:30:1
+   |
+LL | trait D: T + M0 + M1 + M2 {}
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to 4 previous errors
+