feat(tree-sitter): allow Atom-style capture fallbacks (#14196)

theHamsta · web-flow · commit 8ab5ec4aaaee · 2022-02-16T19:38:19.000+01:00
This allows falling back to `@definition` when we have no mapping `@definition.fancy-specialization`. This behavior is described in tree-sitter's documentation (https://tree-sitter.github.io/tree-sitter/syntax-highlighting#theme). Fixes nvim-treesitter/nvim-treesitter#738
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
@@ -22,7 +22,21 @@ local _link_default_highlight_once = function(from, to)
   return from
 end
 
-TSHighlighter.hl_map = {
+-- If @definition.special does not exist use @definition instead
+local subcapture_fallback = {
+  __index = function(self, capture)
+    local rtn
+    local shortened = capture
+    while not rtn and shortened do
+      shortened = shortened:match('(.*)%.')
+      rtn = shortened and rawget(self, shortened)
+    end
+    rawset(self, capture, rtn or "__notfound")
+    return rtn
+  end
+}
+
+TSHighlighter.hl_map = setmetatable({
     ["error"] = "Error",
 
 -- Miscs
@@ -66,7 +80,7 @@ TSHighlighter.hl_map = {
     ["type.builtin"] = "Type",
     ["structure"] = "Structure",
     ["include"] = "Include",
-}
+}, subcapture_fallback)
 
 ---@private
 local function is_highlight_name(capture_name)
diff --git a/test/functional/treesitter/highlight_spec.lua b/test/functional/treesitter/highlight_spec.lua
@@ -23,7 +23,7 @@ local hl_query = [[
   "enum" @type
   "extern" @type
 
-  (string_literal) @string
+  (string_literal) @string.nonexistent-specializer-for-string.should-fallback-to-string
 
   (number_literal) @number
   (char_literal) @string
@@ -613,4 +613,91 @@ describe('treesitter highlighting', function()
       [12] = {background = Screen.colors.Red, bold = true, foreground = Screen.colors.Grey100};
     }}
     end)
+
+  it("allows to use captures with dots (don't use fallback when specialization of foo exists)", function()
+    if pending_c_parser(pending) then return end
+
+    insert([[
+    char* x = "Will somebody ever read this?";
+    ]])
+
+    screen:expect{grid=[[
+      char* x = "Will somebody ever read this?";                       |
+      ^                                                                 |
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+                                                                       |
+    ]]}
+
+    exec_lua [[
+      local parser = vim.treesitter.get_parser(0, "c", {})
+      local highlighter = vim.treesitter.highlighter
+      highlighter.hl_map['foo.bar'] = 'Type'
+      highlighter.hl_map['foo'] = 'String'
+      test_hl = highlighter.new(parser, {queries = {c = "(primitive_type) @foo.bar (string_literal) @foo"}})
+    ]]
+
+    screen:expect{grid=[[
+      {3:char}* x = {5:"Will somebody ever read this?"};                       |
+      ^                                                                 |
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+      {1:~                                                                }|
+                                                                       |
+    ]]}
+  end)
+
+  it("hl_map has the correct fallback behavior", function()
+    exec_lua [[
+      local hl_map = vim.treesitter.highlighter.hl_map
+      hl_map["foo"] = 1
+      hl_map["foo.bar"] = 2
+      hl_map["foo.bar.baz"] = 3
+
+      assert(hl_map["foo"] == 1)
+      assert(hl_map["foo.a.b.c.d"] == 1)
+      assert(hl_map["foo.bar"] == 2)
+      assert(hl_map["foo.bar.a.b.c.d"] == 2)
+      assert(hl_map["foo.bar.baz"] == 3)
+      assert(hl_map["foo.bar.baz.d"] == 3)
+
+      hl_map["FOO"] = 1
+      hl_map["FOO.BAR"] = 2
+      assert(hl_map["FOO.BAR.BAZ"] == 2)
+
+      hl_map["foo.missing.exists"] = 3
+      assert(hl_map["foo.missing"] == 1)
+      assert(hl_map["foo.missing.exists"] == 3)
+      assert(hl_map["foo.missing.exists.bar"] == 3)
+      assert(hl_map["total.nonsense.but.a.lot.of.dots"] == nil)
+      -- It will not perform a second look up of this variable but return a sentinel value
+      assert(hl_map["total.nonsense.but.a.lot.of.dots"] == "__notfound")
+    ]]
+
+  end)
 end)
