1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
use crate::errors::DumpVTableEntries;
use crate::traits::{impossible_predicates, is_vtable_safe_method};
use rustc_hir::def_id::DefId;
use rustc_hir::lang_items::LangItem;
use rustc_infer::traits::util::PredicateSet;
use rustc_infer::traits::ImplSource;
use rustc_middle::query::Providers;
use rustc_middle::traits::BuiltinImplSource;
use rustc_middle::ty::visit::TypeVisitableExt;
use rustc_middle::ty::GenericArgs;
use rustc_middle::ty::{self, GenericParamDefKind, ToPredicate, Ty, TyCtxt, VtblEntry};
use rustc_span::{sym, Span};
use smallvec::SmallVec;

use std::fmt::Debug;
use std::ops::ControlFlow;

#[derive(Clone, Debug)]
pub enum VtblSegment<'tcx> {
    MetadataDSA,
    TraitOwnEntries { trait_ref: ty::PolyTraitRef<'tcx>, emit_vptr: bool },
}

/// Prepare the segments for a vtable
pub fn prepare_vtable_segments<'tcx, T>(
    tcx: TyCtxt<'tcx>,
    trait_ref: ty::PolyTraitRef<'tcx>,
    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
    //    consists of the dsa metadata.
    // 2. Then comes the proper pointer metadata(vptr) and all own methods for all
    //    other super traits except those already included as part of the first
    //    direct super trait virtual table.
    // 3. finally, the own methods of this trait.

    // This has the advantage that trait upcasting to the first direct super trait on each level
    // is zero cost, and to another trait includes only replacing the pointer with one level indirection,
    // while not using too much extra memory.

    // For a single inheritance relationship like this,
    //   D --> C --> B --> A
    // The resulting vtable will consists of these segments:
    //  DSA, A, B, C, D

    // For a multiple inheritance relationship like this,
    //   D --> C --> A
    //           \-> B
    // The resulting vtable will consists of these segments:
    //  DSA, A, B, B-vptr, C, D

    // For a diamond inheritance relationship like this,
    //   D --> B --> A
    //     \-> C -/
    // The resulting vtable will consists of these segments:
    //  DSA, A, B, C, C-vptr, D

    // For a more complex inheritance relationship like this:
    //   O --> G --> C --> A
    //     \     \     \-> B
    //     |     |-> F --> D
    //     |           \-> E
    //     |-> N --> J --> H
    //           \     \-> I
    //           |-> M --> K
    //                 \-> L
    // The resulting vtable will consists of these segments:
    //  DSA, A, B, B-vptr, C, D, D-vptr, E, E-vptr, F, F-vptr, G,
    //  H, H-vptr, I, I-vptr, J, J-vptr, K, K-vptr, L, L-vptr, M, M-vptr,
    //  N, N-vptr, O

    // emit dsa segment first.
    segment_visitor(VtblSegment::MetadataDSA)?;

    let mut emit_vptr_on_new_entry = false;
    let mut visited = PredicateSet::new(tcx);
    let predicate = trait_ref.to_predicate(tcx);
    let mut stack: SmallVec<[(ty::PolyTraitRef<'tcx>, _, _); 5]> =
        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
    // 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
    // stops after it finds a node that has a next-sibling node.
    // This next-sibling node will used as the starting point of next slice.

    // Example:
    // For a diamond inheritance relationship like this,
    //   D#1 --> B#0 --> A#0
    //     \-> C#1 -/

    // Starting point 0 stack [D]
    // Loop run #0: Stack after diving in is [D B A], A is "childless"
    // after this point, all newly visited nodes won't have a vtable that equals to a prefix of this one.
    // Loop run #0: Emitting the slice [B A] (in reverse order), B has a next-sibling node, so this slice stops here.
    // Loop run #0: Stack after exiting out is [D C], C is the next starting point.
    // Loop run #1: Stack after diving in is [D C], C is "childless", since its child A is skipped(already emitted).
    // 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.

    'outer: loop {
        // dive deeper into the stack, recording the path
        'diving_in: loop {
            let &(inner_most_trait_ref, _, _) = stack.last().unwrap();

            let mut direct_super_traits_iter = tcx
                .super_predicates_of(inner_most_trait_ref.def_id())
                .predicates
                .into_iter()
                .filter_map(move |(pred, _)| {
                    pred.instantiate_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,
            }
        }

        // emit innermost item, move to next sibling and stop there if possible, otherwise jump to outer level.
        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: emit_vptr && !tcx.sess.opts.unstable_opts.no_trait_vptr,
            })?;

            // 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;
            }
        }

        // 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,
    trait_ref: ty::PolyTraitRef<'tcx>,
    entries: &[VtblEntry<'tcx>],
) {
    tcx.dcx().emit_err(DumpVTableEntries { span: sp, trait_ref, entries: format!("{entries:#?}") });
}

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);
        let def_id = trait_method.def_id;

        // Some methods cannot be called on an object; skip those.
        if !is_vtable_safe_method(tcx, trait_def_id, trait_method) {
            debug!("own_existential_vtable_entry: not vtable safe");
            return None;
        }

        Some(def_id)
    });

    own_entries
}

/// Given a trait `trait_ref`, iterates the vtable entries
/// that come from `trait_ref`, including its supertraits.
fn vtable_entries<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_ref: ty::PolyTraitRef<'tcx>,
) -> &'tcx [VtblEntry<'tcx>] {
    debug!("vtable_entries({:?})", trait_ref);

    let mut entries = vec![];

    let vtable_segment_callback = |segment| -> ControlFlow<()> {
        match segment {
            VtblSegment::MetadataDSA => {
                entries.extend(TyCtxt::COMMON_VTABLE_ENTRIES);
            }
            VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
                let existential_trait_ref = trait_ref
                    .map_bound(|trait_ref| ty::ExistentialTraitRef::erase_self_ty(tcx, trait_ref));

                // Lookup the shape of vtable for the trait.
                let own_existential_entries =
                    tcx.own_existential_vtable_entries(existential_trait_ref.def_id());

                let own_entries = own_existential_entries.iter().copied().map(|def_id| {
                    debug!("vtable_entries: trait_method={:?}", def_id);

                    // The method may have some early-bound lifetimes; add regions for those.
                    let args = trait_ref.map_bound(|trait_ref| {
                        GenericArgs::for_item(tcx, def_id, |param, _| match param.kind {
                            GenericParamDefKind::Lifetime => tcx.lifetimes.re_erased.into(),
                            GenericParamDefKind::Type { .. }
                            | GenericParamDefKind::Const { .. } => {
                                trait_ref.args[param.index as usize]
                            }
                        })
                    });

                    // The trait type may have higher-ranked lifetimes in it;
                    // erase them if they appear, so that we get the type
                    // at some particular call site.
                    let args =
                        tcx.normalize_erasing_late_bound_regions(ty::ParamEnv::reveal_all(), args);

                    // It's possible that the method relies on where-clauses that
                    // do not hold for this particular set of type parameters.
                    // Note that this method could then never be called, so we
                    // do not want to try and codegen it, in that case (see #23435).
                    let predicates = tcx.predicates_of(def_id).instantiate_own(tcx, args);
                    if impossible_predicates(
                        tcx,
                        predicates.map(|(predicate, _)| predicate).collect(),
                    ) {
                        debug!("vtable_entries: predicates do not hold");
                        return VtblEntry::Vacant;
                    }

                    let instance = ty::Instance::resolve_for_vtable(
                        tcx,
                        ty::ParamEnv::reveal_all(),
                        def_id,
                        args,
                    )
                    .expect("resolution failed during building vtable representation");
                    VtblEntry::Method(instance)
                });

                entries.extend(own_entries);

                if emit_vptr {
                    entries.push(VtblEntry::TraitVPtr(trait_ref));
                }
            }
        }

        ControlFlow::Continue(())
    };

    let _ = prepare_vtable_segments(tcx, trait_ref, vtable_segment_callback);

    if tcx.has_attr(trait_ref.def_id(), sym::rustc_dump_vtable) {
        let sp = tcx.def_span(trait_ref.def_id());
        dump_vtable_entries(tcx, sp, trait_ref, &entries);
    }

    tcx.arena.alloc_from_iter(entries)
}

/// Find slot base for trait methods within vtable entries of another trait
pub(super) fn vtable_trait_first_method_offset<'tcx>(
    tcx: TyCtxt<'tcx>,
    key: (
        ty::PolyTraitRef<'tcx>, // trait_to_be_found
        ty::PolyTraitRef<'tcx>, // trait_owning_vtable
    ),
) -> usize {
    let (trait_to_be_found, trait_owning_vtable) = key;

    // #90177
    let trait_to_be_found_erased = tcx.erase_regions(trait_to_be_found);

    let vtable_segment_callback = {
        let mut vtable_base = 0;

        move |segment| {
            match segment {
                VtblSegment::MetadataDSA => {
                    vtable_base += TyCtxt::COMMON_VTABLE_ENTRIES.len();
                }
                VtblSegment::TraitOwnEntries { trait_ref, emit_vptr } => {
                    if tcx.erase_regions(trait_ref) == trait_to_be_found_erased {
                        return ControlFlow::Break(vtable_base);
                    }
                    vtable_base += count_own_vtable_entries(tcx, trait_ref);
                    if emit_vptr {
                        vtable_base += 1;
                    }
                }
            }
            ControlFlow::Continue(())
        }
    };

    if let Some(vtable_base) =
        prepare_vtable_segments(tcx, trait_owning_vtable, vtable_segment_callback)
    {
        vtable_base
    } else {
        bug!("Failed to find info for expected trait in vtable");
    }
}

/// Find slot offset for trait vptr within vtable entries of another trait
pub(crate) fn vtable_trait_upcasting_coercion_new_vptr_slot<'tcx>(
    tcx: TyCtxt<'tcx>,
    key: (
        Ty<'tcx>, // trait object type whose trait owning vtable
        Ty<'tcx>, // trait object for supertrait
    ),
) -> Option<usize> {
    let (source, target) = key;
    assert!(matches!(&source.kind(), &ty::Dynamic(..)) && !source.has_infer());
    assert!(matches!(&target.kind(), &ty::Dynamic(..)) && !target.has_infer());

    // this has been typecked-before, so diagnostics is not really needed.
    let unsize_trait_did = tcx.require_lang_item(LangItem::Unsize, None);

    let trait_ref = ty::TraitRef::new(tcx, unsize_trait_did, [source, target]);

    match tcx.codegen_select_candidate((ty::ParamEnv::reveal_all(), trait_ref)) {
        Ok(ImplSource::Builtin(BuiltinImplSource::TraitUpcasting { vtable_vptr_slot }, _)) => {
            *vtable_vptr_slot
        }
        otherwise => bug!("expected TraitUpcasting candidate, got {otherwise:?}"),
    }
}

/// Given a trait `trait_ref`, returns the number of vtable entries
/// that come from `trait_ref`, excluding its supertraits. Used in
/// computing the vtable base for an upcast trait of a trait object.
pub(crate) fn count_own_vtable_entries<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_ref: ty::PolyTraitRef<'tcx>,
) -> usize {
    tcx.own_existential_vtable_entries(trait_ref.def_id()).len()
}

pub(super) fn provide(providers: &mut Providers) {
    *providers = Providers {
        own_existential_vtable_entries,
        vtable_entries,
        vtable_trait_upcasting_coercion_new_vptr_slot,
        ..*providers
    };
}