agentic_codebase/index/
language_index.rs

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

agentic_codebase/index/
language_index.rs
