From f7d92f663c07b9be35ceea0342c215ca396c48fb Mon Sep 17 00:00:00 2001
From: Thomas Vigouroux <tomvig38@gmail.com>
Date: Thu, 1 Oct 2020 21:38:38 +0200
Subject: feat(languagetree): implement language tree

Allow the LanguageTree to be used as an option for highlighting.

Co-authored-by: Santos Gallegos <stsewd@protonmail.com>
Co-authored-by: Yazdani Kiyan <yazdani.kiyan@protonmail.com>
---
 README.md                            |   6 --
 lua/nvim-treesitter/highlight.lua    |  19 +++--
 lua/nvim-treesitter/languagetree.lua | 156 +++++++++++++++++++++++++++++++++++
 queries/markdown/highlights.scm      |  19 +++++
 queries/markdown/injections.scm      |   6 ++
 5 files changed, 193 insertions(+), 13 deletions(-)
 create mode 100644 lua/nvim-treesitter/languagetree.lua
 create mode 100644 queries/markdown/highlights.scm
 create mode 100644 queries/markdown/injections.scm

diff --git a/README.md b/README.md
index 320918cdb..575fc4161 100644
--- a/README.md
+++ b/README.md
@@ -104,7 +104,6 @@ All modules are disabled by default,
 so you'll need to activate them by putting this in your `init.vim` file:
 
 ```lua
-lua <<EOF
 require'nvim-treesitter.configs'.setup {
   ensure_installed = "all",     -- one of "all", "language", or a list of languages
   highlight = {
@@ -112,7 +111,6 @@ require'nvim-treesitter.configs'.setup {
     disable = { "c", "rust" },  -- list of language that will be disabled
   },
 }
-EOF
 ```
 
 Check [`:h nvim-treesitter-modules`](doc/nvim-treesitter.txt)
@@ -125,7 +123,6 @@ for a list of available modules and its options.
 Consistent syntax highlighting.
 
 ```lua
-lua <<EOF
 require'nvim-treesitter.configs'.setup {
   highlight = {
     enable = true,
@@ -135,7 +132,6 @@ require'nvim-treesitter.configs'.setup {
     },
   },
 }
-EOF
 ```
 
 ## Incremental selection
@@ -143,7 +139,6 @@ EOF
 Incremental selection based on the named nodes from the grammar.
 
 ```lua
-lua <<EOF
 require'nvim-treesitter.configs'.setup {
   incremental_selection = {
     enable = true,
@@ -155,7 +150,6 @@ require'nvim-treesitter.configs'.setup {
     },
   },
 }
-EOF
 ```
 
 # External modules
diff --git a/lua/nvim-treesitter/highlight.lua b/lua/nvim-treesitter/highlight.lua
index ee68a4ad4..7d21bed6b 100644
--- a/lua/nvim-treesitter/highlight.lua
+++ b/lua/nvim-treesitter/highlight.lua
@@ -81,14 +81,19 @@ function M.attach(bufnr, lang)
   local parser = parsers.get_parser(bufnr, lang)
   local config = configs.get_module('highlight')
 
-  for k, v in pairs(config.custom_captures) do
-    hlmap[k] = v
+  if config.use_languagetree then
+    local ltree = require'nvim-treesitter.languagetree'
+    ltree.new(bufnr, lang)
+  else
+    for k, v in pairs(config.custom_captures) do
+      hlmap[k] = v
+    end
+
+    local query = queries.get_query(lang, "highlights")
+    if not query then return end
+
+    M.highlighters[bufnr] = ts.highlighter.new(parser, query)
   end
-
-  local query = queries.get_query(lang, "highlights")
-  if not query then return end
-
-  M.highlighters[bufnr] = ts.highlighter.new(parser, query)
 end
 
 function M.detach(bufnr)
diff --git a/lua/nvim-treesitter/languagetree.lua b/lua/nvim-treesitter/languagetree.lua
new file mode 100644
index 000000000..fe8828154
--- /dev/null
+++ b/lua/nvim-treesitter/languagetree.lua
@@ -0,0 +1,156 @@
+local parsers = require'nvim-treesitter.parsers'
+local queries = require'nvim-treesitter.query'
+local tsutils = require'nvim-treesitter.ts_utils'
+local TSHighlighter = require'vim.treesitter.highlighter'
+
+local ns = vim.api.nvim_create_namespace("treesitter/highlighter")
+
+local LanguageTree = {}
+LanguageTree.__index = LanguageTree
+
+local trees = { }
+
+function LanguageTree.new(bufnr, lang, not_root)
+  local buf
+  if not bufnr or bufnr == 0 then
+    buf = vim.api.nvim_get_current_buf()
+  else
+    buf = bufnr
+  end
+
+  local parser = parsers.get_parser(buf, parsers.ft_to_lang(lang))
+  if not parser then return end
+
+  local query = queries.get_query(lang, "highlights")
+  if not query then return end
+
+  local self = setmetatable(
+    {
+      highlighter = TSHighlighter.new(parser, query),
+      parser = parser,
+      children = {}
+    },
+    LanguageTree)
+
+  if not not_root then
+    trees[buf] = self
+    self.parser:register_cbs{
+      on_bytes = function() self:update() end
+    }
+  end
+
+  -- First setup
+  self:update()
+
+  return self
+end
+
+function LanguageTree:add_child(lang, child)
+  if not vim.tbl_contains(self.children, child) then
+    table.insert(self.children, child)
+  end
+end
+
+function LanguageTree:remove_child(lang)
+  self.children[lang] = nil
+end
+
+function LanguageTree:node_for_range(range)
+  for _, child in pairs(self.children) do
+    if child:contains(range) then
+      return child:node_for_range(range)
+    end
+  end
+
+  if self:contains(range) then
+    return self
+  end
+end
+
+local function range_contains(source, dest)
+  local start_fits = source[1] < dest[1] or (source[1] == dest[1] and source[2] <= dest[2])
+  local end_fits = source[3] > dest[3] or (source[3] == dest[3] and source[4] >= dest[4])
+
+  return start_fits and end_fits
+end
+
+function LanguageTree:contains(range)
+  for _, source in pairs(self.parser:included_ranges()) do
+    if range_contains(source, range) then
+      return true
+    end
+  end
+
+  return false
+end
+
+function LanguageTree:update()
+  local query = queries.get_query(self.parser.lang, "injections")
+  if not query then return end
+
+  local root = self.parser:parse():root()
+  local startl, _, stopl, _ = root:range()
+
+  local injections = {}
+
+  -- Find injections
+  for inj in queries.iter_prepared_matches(query, root, self.parser.bufnr, startl, stopl+1) do
+    local lang = inj.lang
+
+    if type(lang) ~= "string" then
+      lang = tsutils.get_node_text(lang.node, self.parser.bufnr)[1]
+    end
+
+    if not lang or not inj.injection.node then
+      vim.api.nvim_err_writeln("Invalid match encountered")
+      return nil
+    end
+
+    if not injections[lang] then
+      injections[lang] = {}
+    end
+
+    table.insert(injections[lang], inj.injection.node)
+  end
+
+  local seen = {}
+
+  -- Update each child accordingly
+  for lang, ranges in pairs(injections) do
+    if not self.children[lang] then
+      self.children[lang] = LanguageTree.new(self.parser.bufnr, lang, true)
+    end
+
+    if self.children[lang] then
+      self.children[lang].parser:set_included_ranges(ranges)
+      self.children[lang]:update()
+      seen[lang] = true
+    end
+  end
+
+  -- Clean up unused parsers
+  for lang, _ in pairs(self.children) do
+    if not seen[lang] then
+      self:remove_child(lang)
+    end
+  end
+end
+
+function LanguageTree._on_line(_, _win, buf, line)
+  local tree = trees[buf]
+  if not tree then return end
+
+  local line_len = #(vim.api.nvim_buf_get_lines(buf, line, line + 1, false)[1])
+
+  local matching = tree:node_for_range { line, 0, line, line_len } -- TODO proper search here
+
+  TSHighlighter._on_line("line", _win, buf, line, matching.highlighter)
+end
+
+vim.api.nvim_set_decoration_provider(ns, {
+  on_buf = TSHighlighter._on_buf;
+  on_win = TSHighlighter._on_win;
+  on_line = LanguageTree._on_line;
+})
+
+return LanguageTree
diff --git a/queries/markdown/highlights.scm b/queries/markdown/highlights.scm
new file mode 100644
index 000000000..101526bf5
--- /dev/null
+++ b/queries/markdown/highlights.scm
@@ -0,0 +1,19 @@
+(atx_heading) @text.title
+
+[
+  (code_span)
+  (fenced_code_block)
+]@text.literal
+
+(code_fence_content) @none
+
+[
+  (link_text)
+  (image_description)
+] @text.strong
+
+[
+  (emphasis)
+  (strong_emphasis)
+] @text.emphasis
+(link_destination) @text.uri
diff --git a/queries/markdown/injections.scm b/queries/markdown/injections.scm
new file mode 100644
index 000000000..0522dee2c
--- /dev/null
+++ b/queries/markdown/injections.scm
@@ -0,0 +1,6 @@
+(fenced_code_block
+  (info_string) @lang
+  (code_fence_content) @injection)
+
+((html_block) @injection
+  (#set! "lang" "html"))
-- 
cgit v1.2.3-70-g09d2