Skip to main content

infigraph_core/lang/
registry.rs

1use std::collections::HashMap;
2
3use super::LanguagePack;
4
5/// Probe function: given file content, returns true if this pack should handle the file.
6pub type ContentProbe = fn(content: &[u8]) -> bool;
7
8/// Registry that maps file extensions to language packs.
9pub struct LanguageRegistry {
10    by_extension: HashMap<String, usize>,
11    packs: Vec<LanguagePack>,
12    /// Extension → (pack_index, probe_fn) for content-based override.
13    /// When a file matches a primary extension AND has a content probe registered,
14    /// the probe decides whether to override with the probed pack.
15    content_overrides: HashMap<String, (usize, ContentProbe)>,
16}
17
18impl LanguageRegistry {
19    pub fn new() -> Self {
20        Self {
21            by_extension: HashMap::new(),
22            packs: Vec::new(),
23            content_overrides: HashMap::new(),
24        }
25    }
26
27    /// Register a language pack. All its extensions become resolvable.
28    pub fn register(&mut self, pack: LanguagePack) {
29        let idx = self.packs.len();
30        for ext in &pack.extensions {
31            self.by_extension.insert(ext.clone(), idx);
32        }
33        self.packs.push(pack);
34    }
35
36    /// Register a language pack that claims extensions only when a content probe matches.
37    /// The pack's own `extensions` list may be empty — the `probe_extensions` define
38    /// which extensions trigger the probe. If the probe returns true, this pack overrides
39    /// whatever tree-sitter pack would otherwise handle the file.
40    pub fn register_with_content_probe(
41        &mut self,
42        pack: LanguagePack,
43        probe_extensions: &[&str],
44        probe: ContentProbe,
45    ) {
46        let idx = self.packs.len();
47        for ext in &pack.extensions {
48            self.by_extension.insert(ext.clone(), idx);
49        }
50        for ext in probe_extensions {
51            self.content_overrides.insert(ext.to_string(), (idx, probe));
52        }
53        self.packs.push(pack);
54    }
55
56    /// Look up a language pack by file extension (e.g., ".py").
57    pub fn for_extension(&self, ext: &str) -> Option<&LanguagePack> {
58        self.by_extension.get(ext).map(|&idx| &self.packs[idx])
59    }
60
61    /// Look up a language pack by file path.
62    pub fn for_file(&self, path: &str) -> Option<&LanguagePack> {
63        let ext = path.rsplit_once('.').map(|(_, e)| format!(".{e}"))?;
64        self.for_extension(&ext)
65    }
66
67    /// Look up a language pack by file path and content.
68    /// If a content probe is registered for this extension and matches,
69    /// returns the probed pack instead of the primary extension match.
70    pub fn for_file_with_content(&self, path: &str, content: &[u8]) -> Option<&LanguagePack> {
71        let ext = path.rsplit_once('.').map(|(_, e)| format!(".{e}"))?;
72        if let Some(&(idx, probe)) = self.content_overrides.get(&ext) {
73            if probe(content) {
74                return Some(&self.packs[idx]);
75            }
76        }
77        self.for_extension(&ext)
78    }
79
80    pub fn languages(&self) -> impl Iterator<Item = &LanguagePack> {
81        self.packs.iter()
82    }
83}
84
85impl Default for LanguageRegistry {
86    fn default() -> Self {
87        Self::new()
88    }
89}