tenant.openagents/omega

1use std::borrow::Cow;
2
3use anyhow::Context as _;
4use language_core::{LanguageConfig, LanguageQueries, QUERY_FILENAME_PREFIXES};
5use rust_embed::RustEmbed;
6
7#[derive(RustEmbed)]
8#[folder = "src/"]
9#[exclude = "*.rs"]
10struct GrammarDir;
11
12/// Register all built-in native tree-sitter grammars with the provided registration function.
13///
14/// Each grammar is registered as a `(&str, tree_sitter_language::LanguageFn)` pair.
15/// This must be called before loading language configs/queries.
16#[cfg(feature = "load-grammars")]
17pub fn native_grammars() -> Vec<(&'static str, tree_sitter::Language)> {
18    vec![
19        ("bash", tree_sitter_bash::LANGUAGE.into()),
20        ("c", tree_sitter_c::LANGUAGE.into()),
21        ("cpp", tree_sitter_cpp::LANGUAGE.into()),
22        ("css", tree_sitter_css::LANGUAGE.into()),
23        ("diff", tree_sitter_diff::LANGUAGE.into()),
24        ("go", tree_sitter_go::LANGUAGE.into()),
25        ("gomod", tree_sitter_go_mod::LANGUAGE.into()),
26        ("gowork", tree_sitter_gowork::LANGUAGE.into()),
27        ("jsdoc", tree_sitter_jsdoc::LANGUAGE.into()),
28        ("json", tree_sitter_json::LANGUAGE.into()),
29        ("jsonc", tree_sitter_json::LANGUAGE.into()),
30        ("markdown", tree_sitter_md::LANGUAGE.into()),
31        ("markdown-inline", tree_sitter_md::INLINE_LANGUAGE.into()),
32        ("python", tree_sitter_python::LANGUAGE.into()),
33        ("regex", tree_sitter_regex::LANGUAGE.into()),
34        ("rust", tree_sitter_rust::LANGUAGE.into()),
35        ("tsx", tree_sitter_typescript::LANGUAGE_TSX.into()),
36        (
37            "typescript",
38            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
39        ),
40        ("yaml", tree_sitter_yaml::LANGUAGE.into()),
41        ("gitcommit", tree_sitter_gitcommit::LANGUAGE.into()),
42    ]
43}
44
45/// Load and parse the `config.toml` for a given language name.
46pub fn load_config(name: &str) -> LanguageConfig {
47    let config_toml = String::from_utf8(
48        GrammarDir::get(&format!("{}/config.toml", name))
49            .unwrap_or_else(|| panic!("missing config for language {:?}", name))
50            .data
51            .to_vec(),
52    )
53    .unwrap();
54
55    let config: LanguageConfig = ::toml::from_str(&config_toml)
56        .with_context(|| format!("failed to load config.toml for language {name:?}"))
57        .unwrap();
58
59    config
60}
61
62/// Load and parse the `config.toml` for a given language name, stripping fields
63/// that require grammar support when grammars are not loaded.
64pub fn load_config_for_feature(name: &str, grammars_loaded: bool) -> LanguageConfig {
65    let config = load_config(name);
66
67    if grammars_loaded {
68        config
69    } else {
70        LanguageConfig {
71            name: config.name,
72            matcher: config.matcher,
73            jsx_tag_auto_close: config.jsx_tag_auto_close,
74            ..Default::default()
75        }
76    }
77}
78
79/// Get a raw embedded file by path (relative to `src/`).
80///
81/// Returns the file data as bytes, or `None` if the file does not exist.
82pub fn get_file(path: &str) -> Option<rust_embed::EmbeddedFile> {
83    GrammarDir::get(path)
84}
85
86/// Load all `.scm` query files for a given language name into a `LanguageQueries`.
87///
88/// Multiple `.scm` files with the same prefix (e.g. `highlights.scm` and
89/// `highlights_extra.scm`) are concatenated together with their contents appended.
90pub fn load_queries(name: &str) -> LanguageQueries {
91    let mut result = LanguageQueries::default();
92    for path in GrammarDir::iter() {
93        if let Some(remainder) = path.strip_prefix(name).and_then(|p| p.strip_prefix('/')) {
94            if !remainder.ends_with(".scm") {
95                continue;
96            }
97            for (prefix, query) in QUERY_FILENAME_PREFIXES {
98                if remainder.starts_with(prefix) {
99                    let contents = match GrammarDir::get(path.as_ref()).unwrap().data {
100                        Cow::Borrowed(bytes) => Cow::Borrowed(std::str::from_utf8(bytes).unwrap()),
101                        Cow::Owned(bytes) => Cow::Owned(String::from_utf8(bytes).unwrap()),
102                    };
103                    match query(&mut result) {
104                        None => *query(&mut result) = Some(contents),
105                        Some(existing) => existing.to_mut().push_str(contents.as_ref()),
106                    }
107                }
108            }
109        }
110    }
111    result
112}
113