highlighter: support recursive dynamic injections

crates/ui/src/highlighter/highlighter.rs

@@ -1,4 +1,4 @@
-use crate::highlighter::{HighlightTheme, LanguageRegistry};
+use crate::highlighter::{HighlightTheme, LanguageConfig, LanguageRegistry};
 use crate::input::RopeExt;
 
 use anyhow::{anyhow, Context, Result};
@@ -429,67 +429,155 @@ impl SyntaxHighlighter {
         highlights
     }
 
-    /// TODO: Use incremental parsing to handle the injection.
     fn handle_injection(
         &self,
         injection_language: &str,
         node: Node,
     ) -> Vec<(Range<usize>, String)> {
-        // Ensure byte offsets are on char boundaries for UTF-8 safety
+        let mut ancestors = BTreeSet::new();
-        let start_offset = self.text.clip_offset(node.start_byte(), Bias::Left);
+        self.handle_injection_recursive(
-        let end_offset = self.text.clip_offset(node.end_byte(), Bias::Right);
+            injection_language,
+            node.start_byte(),
+            node.end_byte(),
+            0,
+            &mut ancestors,
+        )
+    }
 
-        let mut cache = vec![];
+    fn lookup_injection_language(name: &str) -> Option<LanguageConfig> {
-        let Some(query) = &self.injection_queries.get(injection_language) else {
+        let registry = LanguageRegistry::singleton();
-            return cache;
+        let names = registry.languages();
+        let exact = |candidate: &str| {
+            names.iter().any(|known| known.as_ref() == candidate)
+                .then(|| registry.language(candidate))
+                .flatten()
         };
-
+        let name = name.to_ascii_lowercase();
-        let content = self.text.slice(start_offset..end_offset);
+        if let Some(config) = exact(&name) {
-        if content.len() == 0 {
+            return Some(config);
-            return cache;
-        };
-        // FIXME: Avoid to_string.
-        let content = content.to_string();
-
-        let Some(config) = LanguageRegistry::singleton().language(injection_language) else {
-            return cache;
-        };
-        let mut parser = Parser::new();
-        if parser.set_language(&config.language).is_err() {
-            return cache;
         }
+        let subtype = name.rsplit_once('/').map(|(_, subtype)| subtype)?;
+        let suffix = subtype.rsplit_once('+').map(|(_, suffix)| suffix).unwrap_or(subtype);
+        exact(suffix).or_else(|| exact(&format!("x-{suffix}")))
+    }
 
-        let source = content.as_bytes();
+    fn handle_injection_recursive(
-        let Some(tree) = parser.parse(source, None) else {
+        &self,
-            return cache;
+        injection_language: &str,
-        };
+        start_byte: usize,
-
+        end_byte: usize,
-        let mut query_cursor = QueryCursor::new();
+        depth: usize,
-        let mut matches = query_cursor.matches(query, tree.root_node(), source);
+        ancestors: &mut BTreeSet<String>,
-
+    ) -> Vec<(Range<usize>, String)> {
-        let mut last_end = start_offset;
+        const MAX_INJECTION_DEPTH: usize = 8;
-        while let Some(m) = matches.next() {
+        if depth >= MAX_INJECTION_DEPTH {
-            for cap in m.captures {
+            return vec![];
-                let cap_node = cap.node;
+        }
-
+        let language = injection_language.to_ascii_lowercase();
-                let node_range: Range<usize> =
+        if !ancestors.insert(language.clone()) {
-                    start_offset + cap_node.start_byte()..start_offset + cap_node.end_byte();
+            return vec![];
-
+        }
-                if node_range.start < last_end {
+        let result = (|| {
-                    continue;
+            let Some(config) = Self::lookup_injection_language(&language) else {
-                }
+                return vec![];
-                if node_range.end > end_offset {
+            };
-                    break;
+            let start = self.text.clip_offset(start_byte, Bias::Left);
-                }
+            let end = self.text.clip_offset(end_byte, Bias::Right);
-
+            if start >= end {
-                if let Some(highlight_name) = query.capture_names().get(cap.index as usize) {
+                return vec![];
-                    last_end = node_range.end;
+            }
-                    cache.push((node_range, highlight_name.to_string()));
+            let content = self.text.slice(start..end).to_string();
+            let source = content.as_bytes();
+            let mut parser = Parser::new();
+            if parser.set_language(&config.language).is_err() {
+                return vec![];
+            }
+            let Some(tree) = parser.parse(source, None) else {
+                return vec![];
+            };
+            let Ok(highlights_query) = Query::new(&config.language, &config.highlights) else {
+                return vec![];
+            };
+            let mut styles = vec![];
+            let mut cursor = QueryCursor::new();
+            let mut matches = cursor.matches(&highlights_query, tree.root_node(), source);
+            while let Some(query_match) = matches.next() {
+                for capture in query_match.captures {
+                    let Some(name) = highlights_query.capture_names().get(capture.index as usize)
+                    else {
+                        continue;
+                    };
+                    styles.push((
+                        start + capture.node.start_byte()..start + capture.node.end_byte(),
+                        name.to_string(),
+                    ));
                 }
             }
-        }
+            if !config.injections.is_empty() {
-
+                if let Ok(injections_query) = Query::new(&config.language, &config.injections) {
-        cache
+                    let content_index = injections_query
+                        .capture_names()
+                        .iter()
+                        .position(|name| *name == "injection.content")
+                        .map(|i| i as u32);
+                    let language_index = injections_query
+                        .capture_names()
+                        .iter()
+                        .position(|name| *name == "injection.language")
+                        .map(|i| i as u32);
+                    let mut cursor = QueryCursor::new();
+                    let mut matches = cursor.matches(&injections_query, tree.root_node(), source);
+                    while let Some(query_match) = matches.next() {
+                        let Some(content_index) = content_index else {
+                            continue;
+                        };
+                        let Some(content_node) = query_match
+                            .captures
+                            .iter()
+                            .find(|capture| capture.index == content_index)
+                            .map(|capture| capture.node)
+                        else {
+                            continue;
+                        };
+                        let mut child_language = language_index.and_then(|index| {
+                            query_match
+                                .captures
+                                .iter()
+                                .find(|capture| capture.index == index)
+                                .and_then(|capture| {
+                                    std::str::from_utf8(
+                                        source.get(capture.node.byte_range())?,
+                                    )
+                                    .ok()
+                                    .map(|text| text.trim().to_ascii_lowercase())
+                                })
+                        });
+                        for prop in injections_query.property_settings(query_match.pattern_index) {
+                            if prop.key.as_ref() == "injection.language" && child_language.is_none() {
+                                child_language = prop
+                                    .value
+                                    .as_ref()
+                                    .map(std::convert::AsRef::as_ref)
+                                    .map(|text| text.to_ascii_lowercase());
+                            }
+                        }
+                        let Some(child_language) = child_language else {
+                            continue;
+                        };
+                        styles.extend(self.handle_injection_recursive(
+                            &child_language,
+                            start + content_node.start_byte(),
+                            start + content_node.end_byte(),
+                            depth + 1,
+                            ancestors,
+                        ));
+                    }
+                }
+            }
+            styles
+        })();
+        ancestors.remove(&language);
+        result
     }
 
     /// Ref:
@@ -506,7 +594,7 @@ impl SyntaxHighlighter {
         query_match: &QueryMatch<'a, 'a>,
     ) -> (Option<SharedString>, Option<Node<'a>>, bool) {
         let content_capture_index = self.injection_content_capture_index;
-        // let language_capture_index = self.injection_language_capture_index;
+        let language_capture_index = self.injection_language_capture_index;
 
         let mut language_name: Option<SharedString> = None;
         let mut content_node = None;
@@ -516,6 +604,12 @@ impl SyntaxHighlighter {
             if index == content_capture_index {
                 content_node = Some(capture.node);
             }
+            if index == language_capture_index && language_name.is_none() {
+                let start = capture.node.start_byte();
+                let end = capture.node.end_byte();
+                let text = self.text.slice(start..end).to_string();
+                language_name = Some(SharedString::from(text.trim().to_ascii_lowercase()));
+            }
         }
 
         let mut include_children = false;