From bbdf62f2635791b4576b50aa7642424d70667418 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 2 Feb 2022 17:01:48 -0800
Subject: [PATCH] Introduce Language::highlight_text method

---
 crates/language/src/buffer.rs   | 33 +++--------------------
 crates/language/src/language.rs | 47 ++++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 30 deletions(-)

diff --git a/crates/language/src/buffer.rs b/crates/language/src/buffer.rs
index 635f2973c6..bcc3d1c6c1 100644
--- a/crates/language/src/buffer.rs
+++ b/crates/language/src/buffer.rs
@@ -21,7 +21,6 @@ use similar::{ChangeTag, TextDiff};
 use smol::future::yield_now;
 use std::{
     any::Any,
-    cell::RefCell,
     cmp::{self, Ordering},
     collections::{BTreeMap, HashMap},
     ffi::OsString,
@@ -38,7 +37,7 @@ use sum_tree::TreeMap;
 use text::{operation_queue::OperationQueue, rope::TextDimension};
 pub use text::{Buffer as TextBuffer, Operation as _, *};
 use theme::SyntaxTheme;
-use tree_sitter::{InputEdit, Parser, QueryCursor, Tree};
+use tree_sitter::{InputEdit, QueryCursor, Tree};
 use util::{post_inc, TryFutureExt as _};
 
 #[cfg(any(test, feature = "test-support"))]
@@ -46,10 +45,6 @@ pub use tree_sitter_rust;
 
 pub use lsp::DiagnosticSeverity;
 
-thread_local! {
-    static PARSER: RefCell<Parser> = RefCell::new(Parser::new());
-}
-
 lazy_static! {
     static ref QUERY_CURSORS: Mutex<Vec<QueryCursor>> = Default::default();
 }
@@ -351,7 +346,7 @@ struct IndentSuggestion {
     indent: bool,
 }
 
-struct TextProvider<'a>(&'a Rope);
+pub(crate) struct TextProvider<'a>(pub(crate) &'a Rope);
 
 struct BufferChunkHighlights<'a> {
     captures: tree_sitter::QueryCaptures<'a, 'a, TextProvider<'a>>,
@@ -938,7 +933,7 @@ impl Buffer {
             let parsed_version = self.version();
             let parse_task = cx.background().spawn({
                 let grammar = grammar.clone();
-                async move { Self::parse_text(&text, old_tree, &grammar) }
+                async move { grammar.parse_text(&text, old_tree) }
             });
 
             match cx
@@ -974,26 +969,6 @@ impl Buffer {
         false
     }
 
-    fn parse_text(text: &Rope, old_tree: Option<Tree>, grammar: &Grammar) -> Tree {
-        PARSER.with(|parser| {
-            let mut parser = parser.borrow_mut();
-            parser
-                .set_language(grammar.ts_language)
-                .expect("incompatible grammar");
-            let mut chunks = text.chunks_in_range(0..text.len());
-            let tree = parser
-                .parse_with(
-                    &mut move |offset, _| {
-                        chunks.seek(offset);
-                        chunks.next().unwrap_or("").as_bytes()
-                    },
-                    old_tree.as_ref(),
-                )
-                .unwrap();
-            tree
-        })
-    }
-
     fn interpolate_tree(&self, tree: &mut SyntaxTree) {
         for edit in self.edits_since::<(usize, Point)>(&tree.version) {
             let (bytes, lines) = edit.flatten();
@@ -2414,7 +2389,7 @@ impl<'a> tree_sitter::TextProvider<'a> for TextProvider<'a> {
     }
 }
 
-struct ByteChunks<'a>(rope::Chunks<'a>);
+pub(crate) struct ByteChunks<'a>(rope::Chunks<'a>);
 
 impl<'a> Iterator for ByteChunks<'a> {
     type Item = &'a [u8];
diff --git a/crates/language/src/language.rs b/crates/language/src/language.rs
index b694e82df1..1dc35e0efa 100644
--- a/crates/language/src/language.rs
+++ b/crates/language/src/language.rs
@@ -17,11 +17,15 @@ use lazy_static::lazy_static;
 pub use outline::{Outline, OutlineItem};
 use parking_lot::Mutex;
 use serde::Deserialize;
-use std::{ops::Range, path::Path, str, sync::Arc};
+use std::{cell::RefCell, ops::Range, path::Path, str, sync::Arc};
 use theme::SyntaxTheme;
 use tree_sitter::{self, Query};
 pub use tree_sitter::{Parser, Tree};
 
+thread_local! {
+    static PARSER: RefCell<Parser>  = RefCell::new(Parser::new());
+}
+
 lazy_static! {
     pub static ref PLAIN_TEXT: Arc<Language> = Arc::new(Language::new(
         LanguageConfig {
@@ -255,6 +259,28 @@ impl Language {
             .and_then(|p| p.label_for_completion(completion))
     }
 
+    pub fn highlight_text<'a>(&'a self, text: &'a Rope) -> Vec<(Range<usize>, HighlightId)> {
+        let mut result = Vec::new();
+        if let Some(grammar) = &self.grammar {
+            let tree = grammar.parse_text(text, None);
+            let mut offset = 0;
+            for chunk in BufferChunks::new(
+                text,
+                0..text.len(),
+                Some(&tree),
+                self.grammar.as_ref(),
+                vec![],
+            ) {
+                let end_offset = offset + chunk.text.len();
+                if let Some(highlight_id) = chunk.highlight_id {
+                    result.push((offset..end_offset, highlight_id));
+                }
+                offset = end_offset;
+            }
+        }
+        result
+    }
+
     pub fn brackets(&self) -> &[BracketPair] {
         &self.config.brackets
     }
@@ -268,6 +294,25 @@ impl Language {
 }
 
 impl Grammar {
+    fn parse_text(&self, text: &Rope, old_tree: Option<Tree>) -> Tree {
+        PARSER.with(|parser| {
+            let mut parser = parser.borrow_mut();
+            parser
+                .set_language(self.ts_language)
+                .expect("incompatible grammar");
+            let mut chunks = text.chunks_in_range(0..text.len());
+            parser
+                .parse_with(
+                    &mut move |offset, _| {
+                        chunks.seek(offset);
+                        chunks.next().unwrap_or("").as_bytes()
+                    },
+                    old_tree.as_ref(),
+                )
+                .unwrap()
+        })
+    }
+
     pub fn highlight_map(&self) -> HighlightMap {
         self.highlight_map.lock().clone()
     }