1use crate::HighlightMap;
  2use anyhow::Result;
  3use parking_lot::Mutex;
  4use serde::Deserialize;
  5use std::{path::Path, str, sync::Arc};
  6use theme::SyntaxTheme;
  7use tree_sitter::{Language as Grammar, Query};
  8pub use tree_sitter::{Parser, Tree};
  9
 10#[derive(Default, Deserialize)]
 11pub struct LanguageConfig {
 12    pub name: String,
 13    pub path_suffixes: Vec<String>,
 14}
 15
 16pub struct Language {
 17    pub(crate) config: LanguageConfig,
 18    pub(crate) grammar: Grammar,
 19    pub(crate) highlight_query: Query,
 20    pub(crate) brackets_query: Query,
 21    pub(crate) highlight_map: Mutex<HighlightMap>,
 22}
 23
 24#[derive(Default)]
 25pub struct LanguageRegistry {
 26    languages: Vec<Arc<Language>>,
 27}
 28
 29impl LanguageRegistry {
 30    pub fn new() -> Self {
 31        Self::default()
 32    }
 33
 34    pub fn add(&mut self, language: Arc<Language>) {
 35        self.languages.push(language);
 36    }
 37
 38    pub fn set_theme(&self, theme: &SyntaxTheme) {
 39        for language in &self.languages {
 40            language.set_theme(theme);
 41        }
 42    }
 43
 44    pub fn select_language(&self, path: impl AsRef<Path>) -> Option<&Arc<Language>> {
 45        let path = path.as_ref();
 46        let filename = path.file_name().and_then(|name| name.to_str());
 47        let extension = path.extension().and_then(|name| name.to_str());
 48        let path_suffixes = [extension, filename];
 49        self.languages.iter().find(|language| {
 50            language
 51                .config
 52                .path_suffixes
 53                .iter()
 54                .any(|suffix| path_suffixes.contains(&Some(suffix.as_str())))
 55        })
 56    }
 57}
 58
 59impl Language {
 60    pub fn new(config: LanguageConfig, grammar: Grammar) -> Self {
 61        Self {
 62            config,
 63            brackets_query: Query::new(grammar, "").unwrap(),
 64            highlight_query: Query::new(grammar, "").unwrap(),
 65            grammar,
 66            highlight_map: Default::default(),
 67        }
 68    }
 69
 70    pub fn with_highlights_query(mut self, highlights_query_source: &str) -> Result<Self> {
 71        self.highlight_query = Query::new(self.grammar, highlights_query_source)?;
 72        Ok(self)
 73    }
 74
 75    pub fn with_brackets_query(mut self, brackets_query_source: &str) -> Result<Self> {
 76        self.brackets_query = Query::new(self.grammar, brackets_query_source)?;
 77        Ok(self)
 78    }
 79
 80    pub fn name(&self) -> &str {
 81        self.config.name.as_str()
 82    }
 83
 84    pub fn highlight_map(&self) -> HighlightMap {
 85        self.highlight_map.lock().clone()
 86    }
 87
 88    pub fn set_theme(&self, theme: &SyntaxTheme) {
 89        *self.highlight_map.lock() = HighlightMap::new(self.highlight_query.capture_names(), theme);
 90    }
 91}
 92
 93#[cfg(test)]
 94mod tests {
 95    use super::*;
 96
 97    #[test]
 98    fn test_select_language() {
 99        let grammar = tree_sitter_rust::language();
100        let registry = LanguageRegistry {
101            languages: vec![
102                Arc::new(Language {
103                    config: LanguageConfig {
104                        name: "Rust".to_string(),
105                        path_suffixes: vec!["rs".to_string()],
106                        ..Default::default()
107                    },
108                    grammar,
109                    highlight_query: Query::new(grammar, "").unwrap(),
110                    brackets_query: Query::new(grammar, "").unwrap(),
111                    highlight_map: Default::default(),
112                }),
113                Arc::new(Language {
114                    config: LanguageConfig {
115                        name: "Make".to_string(),
116                        path_suffixes: vec!["Makefile".to_string(), "mk".to_string()],
117                        ..Default::default()
118                    },
119                    grammar,
120                    highlight_query: Query::new(grammar, "").unwrap(),
121                    brackets_query: Query::new(grammar, "").unwrap(),
122                    highlight_map: Default::default(),
123                }),
124            ],
125        };
126
127        // matching file extension
128        assert_eq!(
129            registry.select_language("zed/lib.rs").map(|l| l.name()),
130            Some("Rust")
131        );
132        assert_eq!(
133            registry.select_language("zed/lib.mk").map(|l| l.name()),
134            Some("Make")
135        );
136
137        // matching filename
138        assert_eq!(
139            registry.select_language("zed/Makefile").map(|l| l.name()),
140            Some("Make")
141        );
142
143        // matching suffix that is not the full file extension or filename
144        assert_eq!(registry.select_language("zed/cars").map(|l| l.name()), None);
145        assert_eq!(
146            registry.select_language("zed/a.cars").map(|l| l.name()),
147            None
148        );
149        assert_eq!(registry.select_language("zed/sumk").map(|l| l.name()), None);
150    }
151}