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