From a91a3e5db073df6e95aec50c395ea17355d2c1c3 Mon Sep 17 00:00:00 2001
From: Wilfred Hughes <me@wilfred.me.uk>
Date: Mon, 18 Apr 2022 09:35:09 -0700
Subject: [PATCH] Handle namespaced highlighting queries

---
 CHANGELOG.md                  |   2 +-
 sample_files/compare.expected |   2 +-
 src/tree_sitter_parser.rs     | 105 +++++++++++++---------------------
 3 files changed, 43 insertions(+), 66 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3785248aa7..ddd4229b96 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,7 +15,7 @@ Improved performance in large files when changes are clustered together.
 
 ### Display
 
-Improved syntax highlighting for keywords.
+Improved syntax highlighting.
 
 Tabs are now rendered with eight spaces.
 
diff --git a/sample_files/compare.expected b/sample_files/compare.expected
index 7a98122be4..e3e3d1f17b 100644
--- a/sample_files/compare.expected
+++ b/sample_files/compare.expected
@@ -32,7 +32,7 @@ sample_files/elisp_contiguous_before.el sample_files/elisp_contiguous_after.el
 e3946aef566a707c718edd7a86340566  -
 
 sample_files/elm_before.elm sample_files/elm_after.elm
-351dd2132fe40414d0b46b5b9380f0fb  -
+c1ea9f99a815b2ae5ce2a7d58fb65368  -
 
 sample_files/hack_before.php sample_files/hack_after.php
 b8c51005df7e1eaaeaf738a4353ac88f  -
diff --git a/src/tree_sitter_parser.rs b/src/tree_sitter_parser.rs
index 2d28c2e9d8..897b3850df 100644
--- a/src/tree_sitter_parser.rs
+++ b/src/tree_sitter_parser.rs
@@ -648,74 +648,51 @@ fn tree_highlights(
     src: &str,
     config: &TreeSitterConfig,
 ) -> HighlightedNodeIds {
-    let mut keyword_ish_capture_ids = vec![];
-    // TODO: Use config.highlight_query.capture_names() to find all
-    // the keyword.foo captures.
-    if let Some(idx) = config.highlight_query.capture_index_for_name("keyword") {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config
-        .highlight_query
-        .capture_index_for_name("keyword.function")
-    {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config
-        .highlight_query
-        .capture_index_for_name("keyword.operator")
-    {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config
-        .highlight_query
-        .capture_index_for_name("keyword.return")
-    {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config.highlight_query.capture_index_for_name("operator") {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config.highlight_query.capture_index_for_name("repeat") {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config.highlight_query.capture_index_for_name("constant") {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config.highlight_query.capture_index_for_name("boolean") {
-        keyword_ish_capture_ids.push(idx);
-    }
-    if let Some(idx) = config
-        .highlight_query
-        .capture_index_for_name("constant.builtin")
-    {
-        keyword_ish_capture_ids.push(idx);
-    }
-
+    let mut keyword_ish_capture_ids: Vec<u32> = vec![];
     let mut string_capture_ids = vec![];
-    if let Some(idx) = config.highlight_query.capture_index_for_name("string") {
-        string_capture_ids.push(idx);
-    }
-    if let Some(idx) = config.highlight_query.capture_index_for_name("character") {
-        string_capture_ids.push(idx);
-    }
-
     let mut type_capture_ids = vec![];
-    if let Some(idx) = config.highlight_query.capture_index_for_name("type") {
-        type_capture_ids.push(idx);
-    }
-    if let Some(idx) = config
-        .highlight_query
-        .capture_index_for_name("type.builtin")
-    {
-        type_capture_ids.push(idx);
-    }
-    if let Some(idx) = config.highlight_query.capture_index_for_name("label") {
+
+    // Query names are often written with namespacing, so
+    // highlights.scm might contain @constant or the more specific
+    // @constant.builtin.
+    //
+    // We support e.g. arbitrary @constant.foo so we get the benefit
+    // of all the relevant highlighting queries.
+    let cn = config.highlight_query.capture_names();
+    for (idx, name) in cn.iter().enumerate() {
+        if name == "type"
+            || name.starts_with("type.")
+            || name.starts_with("storage.type.")
+            || name.starts_with("keyword.type.")
+            || name == "tag"
+        {
+            // TODO: this doesn't capture (type_ref) in Elm as that
+            // applies to the parent node.
+            type_capture_ids.push(idx as u32);
+        } else if name == "keyword"
+            || name.starts_with("keyword.")
+            || name == "constant"
+            || name.starts_with("constant.")
+            || name == "operator"
+            || name == "repeat"
+            || name == "boolean"
+        {
+            keyword_ish_capture_ids.push(idx as u32);
+        }
+
+        if name == "string"
+            || name.starts_with("string.")
+            || name == "character"
+            || name.starts_with("character.")
+        {
+            string_capture_ids.push(idx as u32);
+        }
+
         // Rust uses 'label' for lifetimes, and highglighting
         // lifetimes consistently with types seems reasonable.
-        type_capture_ids.push(idx);
-    }
-    if let Some(idx) = config.highlight_query.capture_index_for_name("tag") {
-        type_capture_ids.push(idx);
+        if name == "label" {
+            type_capture_ids.push(idx as u32);
+        }
     }
 
     let mut qc = ts::QueryCursor::new();