Skip to main content

brokk_bifrost_core/analyzer/
capabilities.rs

1use crate::analyzer::code_unit_index::CodeUnitIndex;
2use crate::analyzer::model::{CodeUnit, ImportInfo, ProjectFile};
3use crate::analyzer::pool_memo::{KeyedPoolSafeMemo, PoolSafeMemo};
4use crate::analyzer::query_token::QueryToken;
5use crate::cancellation::CancellationToken;
6use crate::compact_graph::{CompactRows, CompactRowsBuilder};
7use crate::hash::{HashMap, HashSet};
8use std::any::Any;
9use std::collections::{BTreeSet, VecDeque};
10use std::sync::Arc;
11
12use rayon::prelude::*;
13
14pub trait CapabilityProvider: Any {
15    fn as_any(&self) -> &dyn Any;
16}
17
18impl<T: Any> CapabilityProvider for T {
19    fn as_any(&self) -> &dyn Any {
20        self
21    }
22}
23
24pub trait ImportAnalysisProvider: CapabilityProvider + Send + Sync {
25    /// Shared, not owned: every memoizing implementation already stores the set
26    /// behind an `Arc` in its per-file cache, and the hottest consumers only
27    /// read it (a membership test in candidate discovery, a projection to
28    /// source files in the reverse-import index). Returning the `Arc` removes a
29    /// whole-set clone per cache hit and per insert.
30    fn imported_code_units_of(&self, file: &ProjectFile) -> Arc<HashSet<CodeUnit>>;
31    fn referencing_files_of(&self, file: &ProjectFile) -> HashSet<ProjectFile>;
32
33    /// Return import facts for a group of files without requiring each caller
34    /// to hydrate a complete per-file analyzer state. `None` preserves the
35    /// existing file-at-a-time behavior for providers without a bulk read model.
36    fn import_infos_for_files(
37        &self,
38        _files: &[ProjectFile],
39    ) -> Option<HashMap<ProjectFile, Vec<ImportInfo>>> {
40        None
41    }
42
43    /// The file's import facts, which fall through to the import-tier store
44    /// read when nothing retained answers. Crossing that tier is only cheap
45    /// under an open request scope, so the caller proves one is open by
46    /// passing a [`QueryToken`] (issue #2423).
47    fn import_info_of(&self, _token: QueryToken<'_>, _file: &ProjectFile) -> Vec<ImportInfo> {
48        Vec::new()
49    }
50
51    /// Resolve imported source units from already-loaded import facts. Providers
52    /// that cannot do this cheaply return `None` and use `imported_code_units_of`.
53    fn imported_code_units_from_infos(
54        &self,
55        _file: &ProjectFile,
56        _imports: &[ImportInfo],
57    ) -> Option<Arc<HashSet<CodeUnit>>> {
58        None
59    }
60
61    /// Resolve directly imported project files from already-loaded import facts.
62    /// Providers that do not expose file-level edges return `None` and callers
63    /// can derive a conservative approximation from imported declarations.
64    fn imported_files_from_infos(
65        &self,
66        _file: &ProjectFile,
67        _imports: &[ImportInfo],
68    ) -> Option<HashSet<ProjectFile>> {
69        None
70    }
71
72    fn relevant_imports_for(&self, _code_unit: &CodeUnit) -> HashSet<String> {
73        HashSet::default()
74    }
75
76    /// Whether `source_file` can reference a declaration of `target`. A `true`
77    /// is an answer; a `false` is only an absence of evidence, which is why
78    /// callers still run the expansion backstop after one. Providers that can
79    /// prove the negative override [`Self::import_reachability`] instead and
80    /// derive this from it.
81    fn could_import_file(
82        &self,
83        _source_file: &ProjectFile,
84        _imports: &[ImportInfo],
85        _target: &ProjectFile,
86    ) -> bool {
87        false
88    }
89
90    /// Resolve, in one batch, whatever [`Self::could_import_file`] would look
91    /// up once per candidate.
92    ///
93    /// The shared import-graph candidate walk visits every workspace file, so
94    /// a provider that answers each visit with its own store lookup pays that
95    /// lookup once per import statement in the workspace -- 397k to 662k
96    /// `definition_candidates` round trips inside a single `scan_usages`
97    /// query on a 35k-file Rust workspace (#1748). The walk knows its whole
98    /// candidate set before it inspects any of it, so this hook lets the
99    /// provider enumerate its keys up front and collapse them into batched
100    /// reads.
101    ///
102    /// Doing nothing is always correct: the per-candidate path still answers
103    /// exactly as before, just without the shared warm result. Providers whose
104    /// per-candidate answer is already file-local (Python, JS/TS, Go, C++)
105    /// keep the default.
106    fn prefetch_import_targets(
107        &self,
108        _files: &[ProjectFile],
109        _import_infos: Option<&HashMap<ProjectFile, Vec<ImportInfo>>>,
110        _cancellation: &crate::cancellation::CancellationToken,
111    ) {
112    }
113
114    /// The three-valued form of [`Self::could_import_file`], which lets a
115    /// provider that has a completeness proof retire the caller's backstop.
116    ///
117    /// The default preserves the historical contract exactly by mapping the
118    /// bool spelling's `true` to [`ImportReachability::Reaches`] and its
119    /// `false` to [`ImportReachability::Unknown`] -- never to `DoesNotReach`,
120    /// because the bool contract never distinguished "no" from "I did not
121    /// find one".
122    ///
123    /// Bridging runs in this direction only. A provider overrides this method
124    /// and defines `could_import_file` over it, so the two spellings cannot
125    /// drift apart.
126    fn import_reachability(
127        &self,
128        source_file: &ProjectFile,
129        imports: &[ImportInfo],
130        target: &ProjectFile,
131    ) -> ImportReachability {
132        if self.could_import_file(source_file, imports, target) {
133            ImportReachability::Reaches
134        } else {
135            ImportReachability::Unknown
136        }
137    }
138}
139
140/// How completely a provider can answer "can `source_file` reference a
141/// declaration of `target`?".
142///
143/// Candidate discovery asks this per (candidate, target) pair and, when it
144/// gets no positive answer, materializes every declaration the candidate
145/// imports as a recall backstop. For a namespace-import language that
146/// expansion is the whole workspace's top-level types per using directive,
147/// which is the shape that burned #1194. The backstop exists only because the
148/// bool contract could not distinguish a proven "no" from an unproven one.
149///
150/// A provider must return `DoesNotReach` only from a proof that covers every
151/// way its language can name a declaration without importing it. Anything less
152/// is `Unknown`: the cost of an unnecessary expansion is time, the cost of a
153/// wrong `DoesNotReach` is a missing usage.
154#[derive(Debug, Clone, Copy, PartialEq, Eq)]
155pub enum ImportReachability {
156    /// The candidate can reference the target. Accept it without expanding.
157    Reaches,
158    /// The candidate provably cannot. Reject it without expanding.
159    DoesNotReach,
160    /// Undecided. The caller runs its own backstop, exactly as before this
161    /// verdict existed.
162    Unknown,
163}
164
165/// Resolve direct project-file edges from structured import facts. Prefer a
166/// provider's file-level resolver so imports whose target has no declarations
167/// remain visible; otherwise conservatively project resolved declaration
168/// identities back to their source files.
169pub fn resolve_imported_files_from_infos(
170    provider: &dyn ImportAnalysisProvider,
171    file: &ProjectFile,
172    imports: &[ImportInfo],
173) -> HashSet<ProjectFile> {
174    provider
175        .imported_files_from_infos(file, imports)
176        .unwrap_or_else(|| {
177            provider
178                .imported_code_units_from_infos(file, imports)
179                .unwrap_or_else(|| provider.imported_code_units_of(file))
180                .iter()
181                .map(|unit| unit.source().clone())
182                .collect()
183        })
184}
185
186pub fn build_reverse_import_index<F>(
187    files: &[ProjectFile],
188    resolve_imported: F,
189    parallel: bool,
190) -> HashMap<ProjectFile, Arc<HashSet<ProjectFile>>>
191where
192    F: Fn(&ProjectFile) -> Arc<HashSet<CodeUnit>> + Sync,
193{
194    build_reverse_file_index(
195        files,
196        |file| {
197            resolve_imported(file)
198                .iter()
199                .map(|code_unit| code_unit.source().clone())
200                .collect::<Vec<_>>()
201        },
202        parallel,
203    )
204}
205
206pub type ReverseFileIndex = HashMap<ProjectFile, Arc<HashSet<ProjectFile>>>;
207
208pub fn memoized_reverse_import_index<F, Files>(
209    memo: &PoolSafeMemo<ReverseFileIndex>,
210    files: Files,
211    resolve_imported: F,
212) -> Arc<ReverseFileIndex>
213where
214    F: Fn(&ProjectFile) -> Arc<HashSet<CodeUnit>> + Sync + Copy,
215    Files: Fn() -> Vec<ProjectFile> + Copy,
216{
217    memoized_reverse_file_index(memo, files, |file| {
218        resolve_imported(file)
219            .iter()
220            .map(|code_unit| code_unit.source().clone())
221            .collect::<Vec<_>>()
222    })
223}
224
225pub fn memoized_reverse_file_index<F, I, Files>(
226    memo: &PoolSafeMemo<ReverseFileIndex>,
227    files: Files,
228    resolve_targets: F,
229) -> Arc<ReverseFileIndex>
230where
231    F: Fn(&ProjectFile) -> I + Sync + Copy,
232    I: IntoIterator<Item = ProjectFile>,
233    Files: Fn() -> Vec<ProjectFile> + Copy,
234{
235    memo.get_or_build(
236        || {
237            let files = files();
238            build_reverse_file_index(&files, resolve_targets, true)
239        },
240        || {
241            let files = files();
242            build_reverse_file_index(&files, resolve_targets, false)
243        },
244    )
245}
246
247pub fn build_reverse_file_index<F, I>(
248    files: &[ProjectFile],
249    resolve_targets: F,
250    parallel: bool,
251) -> HashMap<ProjectFile, Arc<HashSet<ProjectFile>>>
252where
253    F: Fn(&ProjectFile) -> I + Sync,
254    I: IntoIterator<Item = ProjectFile>,
255{
256    let collect_edges = |file: &ProjectFile| {
257        let source = file.clone();
258        resolve_targets(file)
259            .into_iter()
260            .filter_map(move |target| (target != source).then(|| (target, source.clone())))
261            .collect::<Vec<_>>()
262    };
263    let edges: Vec<_> = if parallel {
264        files.par_iter().flat_map(collect_edges).collect()
265    } else {
266        files.iter().flat_map(collect_edges).collect()
267    };
268
269    let mut reverse: HashMap<ProjectFile, HashSet<ProjectFile>> = HashMap::default();
270    for (target, source) in edges {
271        reverse.entry(target).or_default().insert(source);
272    }
273    reverse
274        .into_iter()
275        .map(|(file, refs)| (file, Arc::new(refs)))
276        .collect()
277}
278
279/// `Send + Sync` like its three sibling capabilities: every implementor is an
280/// analyzer, and a parallel whole-workspace scan holds the provider across its
281/// fan-out (`kotlin_graph`'s edge builder does exactly that).
282pub trait TypeAliasProvider: CapabilityProvider + Send + Sync {
283    fn is_type_alias(&self, _code_unit: &CodeUnit) -> bool {
284        false
285    }
286}
287
288pub trait TestDetectionProvider: CapabilityProvider {}
289
290/// Which slice of the workspace a descendant index covers.
291///
292/// This is the memo key, and two values is the whole range: the excluded set is
293/// a pure function of the analyzer and the file, so every request that asks to
294/// leave test files out describes the same index.
295#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
296pub enum DescendantIndexVariant {
297    /// Every class-like declaration the analyzer indexes.
298    WholeWorkspace,
299    /// The declarations whose source file the caller's predicate rejected are
300    /// left out. In practice that is a `scan_usages` request with
301    /// `include_tests: false`, whose test files the answer would discard anyway.
302    ProductionOnly,
303}
304
305/// The request state a whole-workspace descendant-index build needs: the
306/// deadline it polls, and the slice of the workspace it covers.
307///
308/// The two travel together because one loop consumes both -- the per-class loop
309/// in [`build_direct_descendant_index_from_candidates`] and, for C++, the
310/// per-file include-closure walk beneath it -- and because the slice is the
311/// memo key that the deadline's complete-or-nothing rule is applied to.
312///
313/// The exclusion predicate is supplied by the caller rather than named here.
314/// Test classification needs an analyzer (`file_is_test_only` reads a language
315/// module index), which this crate must not know about; the closure crossing is
316/// the same one the Rust usage walks already use for their `keep_going`
317/// predicate. **The predicate must be a pure function of the analyzer and the
318/// file.** That is the condition that makes [`DescendantIndexVariant`] a
319/// complete key: two requests that both exclude test files must describe the
320/// same index, or one would be served the other's.
321#[derive(Clone, Copy)]
322pub struct DescendantIndexScope<'a> {
323    cancellation: &'a CancellationToken,
324    excluded_source: Option<&'a dyn Fn(&ProjectFile) -> bool>,
325}
326
327impl<'a> DescendantIndexScope<'a> {
328    /// Every class in the workspace, stopping when `cancellation` says to.
329    pub fn whole_workspace(cancellation: &'a CancellationToken) -> Self {
330        Self {
331            cancellation,
332            excluded_source: None,
333        }
334    }
335
336    /// Every class whose source file `excluded` rejects is left out of the
337    /// index entirely -- it is never handed to `get_direct_ancestors`, so the
338    /// per-file resolution work behind that call is never charged.
339    pub fn excluding_sources(
340        cancellation: &'a CancellationToken,
341        excluded: &'a dyn Fn(&ProjectFile) -> bool,
342    ) -> Self {
343        Self {
344            cancellation,
345            excluded_source: Some(excluded),
346        }
347    }
348
349    pub fn cancellation(&self) -> &'a CancellationToken {
350        self.cancellation
351    }
352
353    pub fn variant(&self) -> DescendantIndexVariant {
354        match self.excluded_source {
355            Some(_) => DescendantIndexVariant::ProductionOnly,
356            None => DescendantIndexVariant::WholeWorkspace,
357        }
358    }
359
360    /// The poll predicate for a builder that would rather take a plain
361    /// `Fn() -> bool` than a token.
362    pub fn keep_going(&self) -> impl Fn() -> bool + use<'_> {
363        || !self.cancellation.is_cancelled()
364    }
365
366    /// Whether `declaration` belongs in this index.
367    pub fn admits(&self, declaration: &CodeUnit) -> bool {
368        self.excluded_source
369            .is_none_or(|excluded| !excluded(declaration.source()))
370    }
371}
372
373pub trait TypeHierarchyProvider: CapabilityProvider + Send + Sync {
374    fn get_direct_ancestors(&self, code_unit: &CodeUnit) -> Vec<CodeUnit>;
375    fn get_direct_descendants(&self, code_unit: &CodeUnit) -> HashSet<CodeUnit>;
376
377    fn supports_type_hierarchy(&self, code_unit: &CodeUnit) -> bool {
378        code_unit.is_class()
379    }
380
381    fn get_ancestors(&self, code_unit: &CodeUnit) -> Vec<CodeUnit> {
382        traverse_hierarchy(code_unit, |next| self.get_direct_ancestors(next))
383    }
384
385    fn get_descendants(&self, code_unit: &CodeUnit) -> Vec<CodeUnit> {
386        traverse_hierarchy(code_unit, |next| {
387            self.get_direct_descendants(next).into_iter().collect()
388        })
389    }
390
391    /// [`Self::get_direct_ancestors`] under a caller's deadline.
392    ///
393    /// `None` means the resolution stopped short; it is never an empty answer,
394    /// which stays the honest way to say "no resolvable base types". The
395    /// default checks the deadline once and delegates, which is right for a
396    /// provider whose ancestor resolution is a lookup. C++ overrides it because
397    /// resolving one class's bases means walking that file's whole transitive
398    /// `#include` closure.
399    fn get_direct_ancestors_within(
400        &self,
401        code_unit: &CodeUnit,
402        scope: &DescendantIndexScope<'_>,
403    ) -> Option<Vec<CodeUnit>> {
404        (!scope.cancellation().is_cancelled()).then(|| self.get_direct_ancestors(code_unit))
405    }
406
407    /// [`Self::get_direct_descendants`] under a caller's deadline and workspace
408    /// slice.
409    ///
410    /// `None` means the answer is not available because the build stopped
411    /// short. It is never a partial answer, and a stopped build publishes
412    /// nothing: a truncated index memoized as complete would be served to every
413    /// later caller as the truth (the rule
414    /// `cancelled_cold_candidate_discovery_does_not_publish_partial_index`
415    /// pins for the Rust walk caches).
416    ///
417    /// The default checks the deadline once and delegates, which is the correct
418    /// behaviour for a provider whose descendant answer is bounded by
419    /// construction. Every provider that inverts the ancestor relation over the
420    /// whole workspace overrides this; a provider that also honours
421    /// [`DescendantIndexScope::admits`] builds one index per
422    /// [`DescendantIndexVariant`]. Ignoring `admits` is sound -- the index is
423    /// then a superset and the caller's own file filter still applies -- so
424    /// only the providers where the prune pays for itself implement it.
425    fn get_direct_descendants_within(
426        &self,
427        code_unit: &CodeUnit,
428        scope: &DescendantIndexScope<'_>,
429    ) -> Option<HashSet<CodeUnit>> {
430        (!scope.cancellation().is_cancelled()).then(|| self.get_direct_descendants(code_unit))
431    }
432
433    /// [`Self::get_descendants`] under a caller's deadline and workspace slice,
434    /// polling once per node of the walk. `None` carries the same meaning as in
435    /// [`Self::get_direct_descendants_within`].
436    fn get_descendants_within(
437        &self,
438        code_unit: &CodeUnit,
439        scope: &DescendantIndexScope<'_>,
440    ) -> Option<Vec<CodeUnit>> {
441        traverse_hierarchy_while(code_unit, scope.cancellation(), |next| {
442            self.get_direct_descendants_within(next, scope)
443                .map(|descendants| descendants.into_iter().collect())
444        })
445    }
446
447    fn get_polymorphic_matches<T: CodeUnitIndex>(
448        &self,
449        target: &CodeUnit,
450        analyzer: &T,
451    ) -> Vec<CodeUnit>
452    where
453        Self: Sized,
454    {
455        if !target.is_function() {
456            return Vec::new();
457        }
458
459        let Some(parent) = analyzer.parent_of(target) else {
460            return Vec::new();
461        };
462
463        self.get_descendants(&parent)
464    }
465}
466
467/// Exact declaration identities plus compact ancestor-to-descendant rows.
468pub struct DirectDescendantIndex {
469    nodes: Box<[CodeUnit]>,
470    row_by_ancestor: HashMap<CodeUnit, u32>,
471    descendants: CompactRows<u32>,
472}
473
474impl DirectDescendantIndex {
475    pub fn from_indexed_nodes(
476        nodes: Vec<CodeUnit>,
477        index_by_node: HashMap<CodeUnit, u32>,
478        mut edges: Vec<(u32, u32)>,
479    ) -> Self {
480        assert_eq!(nodes.len(), index_by_node.len());
481        assert!(nodes.iter().enumerate().all(|(index, node)| {
482            index_by_node.get(node).copied()
483                == Some(
484                    u32::try_from(index).expect("hierarchy index declarations must fit in a u32"),
485                )
486        }));
487        assert!(edges.iter().all(|(ancestor, descendant)| {
488            (*ancestor as usize) < nodes.len() && (*descendant as usize) < nodes.len()
489        }));
490        edges.sort_unstable();
491        edges.dedup();
492
493        let row_count = usize::from(!edges.is_empty())
494            + edges
495                .windows(2)
496                .filter(|pair| pair[0].0 != pair[1].0)
497                .count();
498        let mut row_by_ancestor = HashMap::default();
499        let mut descendants = CompactRowsBuilder::with_capacity(row_count, edges.len());
500        let mut cursor = 0;
501        while cursor < edges.len() {
502            let ancestor = edges[cursor].0;
503            let start = cursor;
504            while cursor < edges.len() && edges[cursor].0 == ancestor {
505                cursor += 1;
506            }
507            let row =
508                u32::try_from(descendants.rows()).expect("hierarchy index rows must fit in a u32");
509            row_by_ancestor.insert(nodes[ancestor as usize].clone(), row);
510            descendants.push_row(
511                edges[start..cursor]
512                    .iter()
513                    .map(|(_, descendant)| *descendant),
514            );
515        }
516        Self {
517            nodes: nodes.into_boxed_slice(),
518            row_by_ancestor,
519            descendants: descendants.finish(),
520        }
521    }
522
523    pub fn descendants(&self, ancestor: &CodeUnit) -> HashSet<CodeUnit> {
524        let Some(row) = self.row_by_ancestor.get(ancestor).copied() else {
525            return HashSet::default();
526        };
527        self.descendants
528            .row(row as usize)
529            .iter()
530            .map(|descendant| self.nodes[*descendant as usize].clone())
531            .collect()
532    }
533}
534
535/// Answer a descendant query from a variant-keyed index cell, building the
536/// variant the caller asked for if it is missing.
537///
538/// Every analyzer that memoizes a whole-workspace descendant index repeats this
539/// same dance: pick the cell for the scope's variant, build while the deadline
540/// holds, publish nothing if it does not, project the ancestor's row. It lives
541/// here so the complete-or-nothing rule is stated once instead of copied into
542/// each language module.
543pub fn descendants_from_variant_index(
544    index: &KeyedPoolSafeMemo<DescendantIndexVariant, DirectDescendantIndex>,
545    scope: &DescendantIndexScope<'_>,
546    code_unit: &CodeUnit,
547    build: impl Fn() -> Option<DirectDescendantIndex>,
548) -> Option<HashSet<CodeUnit>> {
549    Some(
550        index
551            .cell(&scope.variant())
552            // The builders are serial, so the same closure serves both memo
553            // arms; the memo's value here is the non-blocking claim protocol.
554            .get_or_build_while(&scope.keep_going(), &build, &build)?
555            .descendants(code_unit),
556    )
557}
558
559/// Invert the ancestor relation over every class the analyzer indexes.
560///
561/// `scope` decides two things. Its predicate drops declarations before they are
562/// ever handed to `get_direct_ancestors`, which is where the per-declaration
563/// resolution cost is charged (for C++ that is a whole transitive `#include`
564/// closure per declaring file). Its token bounds the loop: `None` means the
565/// build stopped short and must not be published.
566pub fn build_direct_descendant_index<A, P>(
567    analyzer: &A,
568    provider: &P,
569    scope: &DescendantIndexScope<'_>,
570) -> Option<DirectDescendantIndex>
571where
572    A: CodeUnitIndex,
573    P: TypeHierarchyProvider + ?Sized,
574{
575    build_direct_descendant_index_from_candidates(
576        analyzer
577            .all_declarations()
578            .filter(|candidate| candidate.is_class() && scope.admits(candidate))
579            .collect(),
580        |candidate| provider.get_direct_ancestors_within(candidate, scope),
581        &scope.keep_going(),
582    )
583}
584
585/// The edge-building half of [`build_direct_descendant_index`], for providers
586/// that assemble their candidate list some other way.
587///
588/// `keep_going` is polled once per candidate. That granularity is the natural
589/// checkpoint: one candidate is one `direct_ancestors` call, which is the unit
590/// of work whose tens of thousands of repetitions made this loop unbounded in
591/// the first place (issue #1748). `direct_ancestors` answers `None` when its
592/// own work stopped short, which stops this loop too -- a candidate whose bases
593/// were half-resolved must not contribute a half-populated edge set.
594pub fn build_direct_descendant_index_from_candidates<F>(
595    mut candidates: Vec<CodeUnit>,
596    mut direct_ancestors: F,
597    keep_going: &dyn Fn() -> bool,
598) -> Option<DirectDescendantIndex>
599where
600    F: FnMut(&CodeUnit) -> Option<Vec<CodeUnit>>,
601{
602    candidates.sort();
603    candidates.dedup();
604    let mut types_by_fq_name: HashMap<String, Vec<CodeUnit>> = HashMap::default();
605    for candidate in &candidates {
606        types_by_fq_name
607            .entry(candidate.fq_name())
608            .or_default()
609            .push(candidate.clone());
610    }
611    let mut nodes = candidates.clone();
612    let mut index_by_node: HashMap<_, _> = nodes
613        .iter()
614        .enumerate()
615        .map(|(index, node)| {
616            (
617                node.clone(),
618                u32::try_from(index).expect("hierarchy index declarations must fit in a u32"),
619            )
620        })
621        .collect();
622    let mut edges = Vec::new();
623    for candidate in candidates {
624        if !keep_going() {
625            return None;
626        }
627        let descendant = index_by_node[&candidate];
628        for ancestor in direct_ancestors(&candidate)? {
629            let ancestor = types_by_fq_name
630                .get(&ancestor.fq_name())
631                .and_then(|same_name| {
632                    let mut same_source = same_name
633                        .iter()
634                        .filter(|unit| unit.source() == candidate.source());
635                    let exact = same_source.next()?;
636                    same_source.next().is_none().then(|| exact.clone())
637                })
638                .unwrap_or(ancestor);
639            let ancestor = *index_by_node.entry(ancestor.clone()).or_insert_with(|| {
640                let index = u32::try_from(nodes.len())
641                    .expect("hierarchy index declarations must fit in a u32");
642                nodes.push(ancestor);
643                index
644            });
645            edges.push((ancestor, descendant));
646        }
647    }
648    Some(DirectDescendantIndex::from_indexed_nodes(
649        nodes,
650        index_by_node,
651        edges,
652    ))
653}
654
655/// [`traverse_hierarchy`] under a deadline: polls once per node popped and
656/// propagates a stopped step as `None` rather than returning the nodes it had
657/// already collected. A partial subtype set read as complete would let a caller
658/// conclude that a class has no further subclasses.
659fn traverse_hierarchy_while<F>(
660    root: &CodeUnit,
661    cancellation: &CancellationToken,
662    mut next: F,
663) -> Option<Vec<CodeUnit>>
664where
665    F: FnMut(&CodeUnit) -> Option<Vec<CodeUnit>>,
666{
667    let direct = next(root)?;
668    if direct.is_empty() {
669        return Some(Vec::new());
670    }
671
672    let mut seen = BTreeSet::new();
673    let mut result = Vec::new();
674    let mut queue = VecDeque::new();
675
676    for item in direct {
677        if seen.insert(item.fq_name()) {
678            queue.push_back(item.clone());
679            result.push(item);
680        }
681    }
682
683    while let Some(current) = queue.pop_front() {
684        if cancellation.is_cancelled() {
685            return None;
686        }
687        for item in next(&current)? {
688            if seen.insert(item.fq_name()) {
689                queue.push_back(item.clone());
690                result.push(item);
691            }
692        }
693    }
694
695    Some(result)
696}
697
698fn traverse_hierarchy<F>(root: &CodeUnit, mut next: F) -> Vec<CodeUnit>
699where
700    F: FnMut(&CodeUnit) -> Vec<CodeUnit>,
701{
702    let direct = next(root);
703    if direct.is_empty() {
704        return Vec::new();
705    }
706
707    let mut seen = BTreeSet::new();
708    let mut result = Vec::new();
709    let mut queue = VecDeque::new();
710
711    for item in direct {
712        if seen.insert(item.fq_name()) {
713            queue.push_back(item.clone());
714            result.push(item);
715        }
716    }
717
718    while let Some(current) = queue.pop_front() {
719        for item in next(&current) {
720            if seen.insert(item.fq_name()) {
721                queue.push_back(item.clone());
722                result.push(item);
723            }
724        }
725    }
726
727    result
728}