agentic_codebase/index/
type_index.rs

1//! Index by code unit type.
2//!
3//! Groups code units by their [`CodeUnitType`] for fast filtered lookups.
4//! Building the index is O(n) in the number of units; lookups are O(1).
5
6use std::collections::HashMap;
7
8use crate::graph::CodeGraph;
9use crate::types::CodeUnitType;
10
11/// Index that groups code units by their type.
12#[derive(Debug, Clone, Default)]
13pub struct TypeIndex {
14    by_type: HashMap<CodeUnitType, Vec<u64>>,
15}
16
17impl TypeIndex {
18    /// Build a `TypeIndex` from all code units in the given graph.
19    pub fn build(graph: &CodeGraph) -> Self {
20        let mut by_type: HashMap<CodeUnitType, Vec<u64>> = HashMap::new();
21
22        for unit in graph.units() {
23            by_type.entry(unit.unit_type).or_default().push(unit.id);
24        }
25
26        Self { by_type }
27    }
28
29    /// Look up all unit IDs of the given type.
30    ///
31    /// Returns an empty slice if no units match.
32    pub fn lookup(&self, unit_type: CodeUnitType) -> &[u64] {
33        self.by_type
34            .get(&unit_type)
35            .map(|v| v.as_slice())
36            .unwrap_or(&[])
37    }
38
39    /// Returns the count of units of the given type.
40    pub fn count(&self, unit_type: CodeUnitType) -> usize {
41        self.by_type.get(&unit_type).map(|v| v.len()).unwrap_or(0)
42    }
43
44    /// Returns all code unit types present in the index.
45    pub fn types(&self) -> Vec<CodeUnitType> {
46        self.by_type.keys().copied().collect()
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53    use crate::graph::CodeGraph;
54    use crate::types::{CodeUnit, CodeUnitType, Language, Span};
55    use std::path::PathBuf;
56
57    fn make_unit(unit_type: CodeUnitType) -> CodeUnit {
58        CodeUnit::new(
59            unit_type,
60            Language::Rust,
61            "test_unit".to_string(),
62            "mod::test_unit".to_string(),
63            PathBuf::from("src/lib.rs"),
64            Span::new(1, 0, 10, 0),
65        )
66    }
67
68    #[test]
69    fn test_empty_index() {
70        let graph = CodeGraph::default();
71        let index = TypeIndex::build(&graph);
72        assert_eq!(index.count(CodeUnitType::Function), 0);
73        assert_eq!(index.lookup(CodeUnitType::Function), &[] as &[u64]);
74        assert!(index.types().is_empty());
75    }
76
77    #[test]
78    fn test_grouped_lookup() {
79        let mut graph = CodeGraph::default();
80        graph.add_unit(make_unit(CodeUnitType::Function));
81        graph.add_unit(make_unit(CodeUnitType::Function));
82        graph.add_unit(make_unit(CodeUnitType::Module));
83        graph.add_unit(make_unit(CodeUnitType::Type));
84
85        let index = TypeIndex::build(&graph);
86        assert_eq!(index.count(CodeUnitType::Function), 2);
87        assert_eq!(index.count(CodeUnitType::Module), 1);
88        assert_eq!(index.count(CodeUnitType::Type), 1);
89        assert_eq!(index.count(CodeUnitType::Import), 0);
90
91        assert_eq!(index.lookup(CodeUnitType::Function), &[0, 1]);
92        assert_eq!(index.lookup(CodeUnitType::Module), &[2]);
93    }
94
95    #[test]
96    fn test_types_list() {
97        let mut graph = CodeGraph::default();
98        graph.add_unit(make_unit(CodeUnitType::Function));
99        graph.add_unit(make_unit(CodeUnitType::Trait));
100
101        let index = TypeIndex::build(&graph);
102        let mut types = index.types();
103        types.sort_by_key(|t| *t as u8);
104        assert_eq!(types, vec![CodeUnitType::Function, CodeUnitType::Trait]);
105    }
106}
agentic_codebase/index/type_index.rs

agentic_codebase/index/
type_index.rs
