Skip to main content

hearth_graph/
bundled.rs

1use std::borrow::Cow;
2
3use tree_sitter::Language;
4
5use crate::{ImportKind, ImportSpec, LanguageRegistry, LanguageSpec};
6
7const JAVASCRIPT_IMPORTS_QUERY: &str = include_str!("../queries/javascript/imports.scm");
8const TYPESCRIPT_IMPORTS_QUERY: &str = include_str!("../queries/typescript/imports.scm");
9
10fn language_spec(
11    name: &'static str,
12    language: Language,
13    extensions: &[&str],
14    tags_query: Cow<'static, str>,
15) -> LanguageSpec {
16    language_spec_with_imports(name, language, extensions, tags_query, None)
17}
18
19fn language_spec_with_imports(
20    name: &'static str,
21    language: Language,
22    extensions: &[&str],
23    tags_query: Cow<'static, str>,
24    imports: Option<ImportSpec>,
25) -> LanguageSpec {
26    let spec = LanguageSpec::new(name, language, extensions).with_tags_query(tags_query);
27    match imports {
28        Some(imports) => spec.with_imports(imports),
29        None => spec,
30    }
31}
32
33fn language_spec_merging_adjacent_definitions(
34    name: &'static str,
35    language: Language,
36    extensions: &[&str],
37    tags_query: Cow<'static, str>,
38) -> LanguageSpec {
39    language_spec(name, language, extensions, tags_query)
40        .with_merge_adjacent_same_name_definitions(true)
41}
42
43fn import_kind(capture: &str) -> ImportKind {
44    match capture {
45        "import.source.static" => ImportKind::EsStatic,
46        "import.source.reexport" => ImportKind::EsReexport,
47        "import.source.dynamic" => ImportKind::EsDynamic,
48        "import.source.commonjs" => ImportKind::CommonJs,
49        "import.source.tsrequire" => ImportKind::TsImportRequire,
50        _ => unreachable!("unexpected import capture: {capture}"),
51    }
52}
53
54impl LanguageRegistry {
55    /// Creates a registry containing Hearth's bundled language grammars.
56    #[must_use]
57    pub fn bundled() -> Self {
58        let mut registry = Self::empty();
59
60        registry.register(language_spec_with_imports(
61            "rust",
62            tree_sitter_rust::LANGUAGE.into(),
63            &["rs"],
64            Cow::Borrowed(tree_sitter_rust::TAGS_QUERY),
65            Some(ImportSpec::Custom(crate::imports::rust::extract)),
66        ));
67        registry.register(language_spec_with_imports(
68            "typescript",
69            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
70            &["ts", "mts", "cts"],
71            Cow::Owned(format!(
72                "{}\n{}",
73                tree_sitter_javascript::TAGS_QUERY,
74                tree_sitter_typescript::TAGS_QUERY
75            )),
76            Some(ImportSpec::Query {
77                source: Cow::Owned(format!(
78                    "{JAVASCRIPT_IMPORTS_QUERY}\n{TYPESCRIPT_IMPORTS_QUERY}"
79                )),
80                kind_map: import_kind,
81            }),
82        ));
83        registry.register(language_spec_with_imports(
84            "tsx",
85            tree_sitter_typescript::LANGUAGE_TSX.into(),
86            &["tsx"],
87            Cow::Owned(format!(
88                "{}\n{}",
89                tree_sitter_javascript::TAGS_QUERY,
90                tree_sitter_typescript::TAGS_QUERY
91            )),
92            Some(ImportSpec::Query {
93                source: Cow::Owned(format!(
94                    "{JAVASCRIPT_IMPORTS_QUERY}\n{TYPESCRIPT_IMPORTS_QUERY}"
95                )),
96                kind_map: import_kind,
97            }),
98        ));
99        registry.register(language_spec_with_imports(
100            "javascript",
101            tree_sitter_javascript::LANGUAGE.into(),
102            &["js", "mjs", "cjs"],
103            Cow::Borrowed(tree_sitter_javascript::TAGS_QUERY),
104            Some(ImportSpec::Query {
105                source: Cow::Borrowed(JAVASCRIPT_IMPORTS_QUERY),
106                kind_map: import_kind,
107            }),
108        ));
109        registry.register(language_spec_with_imports(
110            "jsx",
111            tree_sitter_javascript::LANGUAGE.into(),
112            &["jsx"],
113            Cow::Borrowed(tree_sitter_javascript::TAGS_QUERY),
114            Some(ImportSpec::Query {
115                source: Cow::Borrowed(JAVASCRIPT_IMPORTS_QUERY),
116                kind_map: import_kind,
117            }),
118        ));
119        registry.register(language_spec(
120            "go",
121            tree_sitter_go::LANGUAGE.into(),
122            &["go"],
123            Cow::Borrowed(tree_sitter_go::TAGS_QUERY),
124        ));
125        registry.register(language_spec(
126            "python",
127            tree_sitter_python::LANGUAGE.into(),
128            &["py"],
129            Cow::Borrowed(tree_sitter_python::TAGS_QUERY),
130        ));
131        registry.register(language_spec(
132            "ruby",
133            tree_sitter_ruby::LANGUAGE.into(),
134            &["rb", "rake", "gemspec"],
135            Cow::Borrowed(tree_sitter_ruby::TAGS_QUERY),
136        ));
137        registry.register(language_spec(
138            "c",
139            tree_sitter_c::LANGUAGE.into(),
140            &["c", "h"],
141            Cow::Borrowed(tree_sitter_c::TAGS_QUERY),
142        ));
143        registry.register(language_spec(
144            "cpp",
145            tree_sitter_cpp::LANGUAGE.into(),
146            &["cpp", "cc", "cxx", "hpp", "hxx"],
147            Cow::Borrowed(tree_sitter_cpp::TAGS_QUERY),
148        ));
149        registry.register(language_spec(
150            "java",
151            tree_sitter_java::LANGUAGE.into(),
152            &["java"],
153            Cow::Borrowed(tree_sitter_java::TAGS_QUERY),
154        ));
155        registry.register(language_spec(
156            "csharp",
157            tree_sitter_c_sharp::LANGUAGE.into(),
158            &["cs"],
159            Cow::Borrowed(include_str!("../queries/c_sharp/tags.scm")),
160        ));
161        registry.register(language_spec(
162            "zig",
163            tree_sitter_zig::LANGUAGE.into(),
164            &["zig"],
165            Cow::Borrowed(include_str!("../queries/zig/tags.scm")),
166        ));
167        registry.register(language_spec(
168            "bash",
169            tree_sitter_bash::LANGUAGE.into(),
170            &["sh", "bash", "zsh"],
171            Cow::Borrowed(include_str!("../queries/bash/tags.scm")),
172        ));
173        registry.register(language_spec_merging_adjacent_definitions(
174            "haskell",
175            tree_sitter_haskell::LANGUAGE.into(),
176            &["hs", "lhs"],
177            Cow::Borrowed(include_str!("../queries/haskell/tags.scm")),
178        ));
179        registry.register(language_spec(
180            "lua",
181            tree_sitter_lua::LANGUAGE.into(),
182            &["lua"],
183            Cow::Borrowed(tree_sitter_lua::TAGS_QUERY),
184        ));
185        registry.register(language_spec(
186            "php",
187            tree_sitter_php::LANGUAGE_PHP.into(),
188            &["php"],
189            Cow::Borrowed(tree_sitter_php::TAGS_QUERY),
190        ));
191        registry.register(language_spec(
192            "swift",
193            tree_sitter_swift::LANGUAGE.into(),
194            &["swift"],
195            Cow::Borrowed(tree_sitter_swift::TAGS_QUERY),
196        ));
197        registry.register(language_spec(
198            "markdown",
199            tree_sitter_md::LANGUAGE.into(),
200            &["md", "markdown"],
201            Cow::Borrowed(include_str!("../queries/markdown/tags.scm")),
202        ));
203
204        registry
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use std::path::Path;
211
212    use super::*;
213    use crate::ParserPool;
214
215    #[test]
216    fn all_bundled_tags_queries_compile_and_are_pool_reachable() {
217        let registry = LanguageRegistry::bundled();
218        assert_eq!(registry.iter().len(), 19);
219
220        let mut pool = ParserPool::new(&registry);
221        for (id, spec) in registry.iter() {
222            let source = spec
223                .tags_query
224                .as_deref()
225                .unwrap_or_else(|| panic!("{} is missing its tags query", spec.name));
226            tree_sitter::Query::new(&spec.language, source)
227                .unwrap_or_else(|error| panic!("{} tags query failed: {error}", spec.name));
228            assert!(
229                pool.tags_query(id).is_some(),
230                "{} is not reachable through ParserPool::tags_query",
231                spec.name
232            );
233        }
234    }
235
236    #[test]
237    fn all_bundled_import_queries_compile_and_are_pool_reachable() {
238        let registry = LanguageRegistry::bundled();
239        let mut pool = ParserPool::new(&registry);
240        let mut query_count = 0;
241
242        for (id, spec) in registry.iter() {
243            match spec.imports.as_ref() {
244                Some(ImportSpec::Query { source, .. }) => {
245                    query_count += 1;
246                    tree_sitter::Query::new(&spec.language, source).unwrap_or_else(|error| {
247                        panic!("{} imports query failed: {error}", spec.name)
248                    });
249                    assert!(
250                        pool.imports_query(id).is_some(),
251                        "{} is not reachable through ParserPool::imports_query",
252                        spec.name
253                    );
254                }
255                Some(ImportSpec::Custom(_)) | None => {
256                    assert!(pool.imports_query(id).is_none(), "{}", spec.name);
257                }
258            }
259        }
260
261        assert_eq!(query_count, 4);
262    }
263
264    #[test]
265    fn register_uses_last_extension_owner_and_increments_generation() {
266        let mut registry = LanguageRegistry::bundled();
267        let original_id = registry
268            .for_path(Path::new("main.rs"))
269            .expect("bundled Rust extension");
270        let original_generation = registry.generation();
271
272        let replacement_id = registry.register(language_spec(
273            "replacement-rust",
274            tree_sitter_rust::LANGUAGE.into(),
275            &["rs"],
276            Cow::Borrowed(tree_sitter_rust::TAGS_QUERY),
277        ));
278
279        assert_ne!(replacement_id, original_id);
280        assert_eq!(
281            registry.for_path(Path::new("main.rs")),
282            Some(replacement_id)
283        );
284        assert_eq!(registry.generation(), original_generation + 1);
285    }
286
287    #[test]
288    fn supports_symbols_for_bundled_module_extensions() {
289        let registry = LanguageRegistry::bundled();
290
291        for path in [
292            "module.mjs",
293            "module.cjs",
294            "module.mts",
295            "module.cts",
296            "module.ts",
297            "component.tsx",
298            "lib.rs",
299        ] {
300            assert!(registry.supports_symbols(Path::new(path)), "{path}");
301        }
302
303        for path in ["style.css", "component.vue", "Makefile"] {
304            assert!(!registry.supports_symbols(Path::new(path)), "{path}");
305        }
306    }
307
308    #[test]
309    fn supports_imports_for_bundled_module_extensions() {
310        let registry = LanguageRegistry::bundled();
311
312        for path in [
313            "module.mjs",
314            "module.cjs",
315            "module.mts",
316            "module.cts",
317            "module.ts",
318            "component.tsx",
319            "component.jsx",
320            "lib.rs",
321        ] {
322            assert!(registry.supports_imports(Path::new(path)), "{path}");
323        }
324
325        for path in ["main.go", "style.css", "Makefile"] {
326            assert!(!registry.supports_imports(Path::new(path)), "{path}");
327        }
328    }
329}