Treesitter indent

also fixes the memoize_by_buf_tick function

5 files changed, 79 insertions, 55 deletions

diff --git a/lua/nvim-treesitter/configs.lua b/lua/nvim-treesitter/configs.lua
index 2655adc79..63165f9e0 100644
--- a/lua/nvim-treesitter/configs.lua
+++ b/lua/nvim-treesitter/configs.lua
@@ -40,7 +40,7 @@ local builtin_modules = {
     module_path = 'nvim-treesitter.indent',
     enable = false,
     disable = {},
-    is_supported = queries.has_locals
+    is_supported = queries.has_indents
   }
 }
 
diff --git a/lua/nvim-treesitter/indent.lua b/lua/nvim-treesitter/indent.lua
index bd7d64454..b747156af 100644
--- a/lua/nvim-treesitter/indent.lua
+++ b/lua/nvim-treesitter/indent.lua
@@ -1,58 +1,64 @@
--- TODO: logic not working properly
--- need to find a better way to get the indent from a line
--- it almosts works though
-local api = vim.api
-local utils = require'nvim-treesitter.ts_utils'
-local query = require'nvim-treesitter.query'
 local parsers = require'nvim-treesitter.parsers'
+local queries = require'nvim-treesitter.query'
 local utils = require'nvim-treesitter.ts_utils'
-local locals = require'nvim-treesitter.locals'
 
 local M = {}
 
-local function get_node_at_line(lnum)
-  local node = utils.get_node_at_cursor()
-  local srow = node:range()
+local function get_node_at_line(root, lnum)
+  for node in root:iter_children() do
+    local srow, _, erow = node:range()
+    if srow == lnum then return node end
 
-  if srow+1 < lnum then
-    for n in node:iter_children() do
-      local row = n:range()
-      if row+1 == lnum then
-        node = n
-        break
-      end
+    if node:child_count() > 0 and srow < lnum and lnum <= erow then
+      return get_node_at_line(node, lnum)
     end
   end
 
-  return node
+  local wrapper = root:descendant_for_range(lnum, 0, lnum, -1)
+  local child = wrapper:child(0)
+  return child or wrapper
 end
 
-function M.get_indent(lnum)
-  if not parsers.has_parser() or not lnum then return -1 end
+local get_indents = utils.memoize_by_buf_tick(function(bufnr)
+  local indents = queries.get_capture_matches(bufnr, '@indent.node', 'indents') or {}
+  local branches = queries.get_capture_matches(bufnr, '@branch.node', 'indents') or {}
+
+  local indents_map = {}
+  for _, node in ipairs(indents) do
+    indents_map[tostring(node)] = true
+  end
 
-  local node = get_node_at_line(lnum)
-  local srow, scol, erow, ecol = node:range()
+  local branches_map = {}
+  for _, node in ipairs(branches) do
+    branches_map[tostring(node)] = true
+  end
+
+  return { indents = indents_map, branches = branches_map }
+end)
 
-  local parent = locals.containing_scope(node:parent() or node)
-  local parent_row, pscol, perow, pecol = parent:range()
-  local parent_indent = vim.fn.indent(parent_row+1)
+function M.get_indent(lnum)
+  local parser = parsers.get_parser()
+  if not parser or not lnum then return -1 end
 
-  if (parent_row == srow and pscol == scol) or (perow == erow and pecol == ecol) then
-    node = parent
-    srow = node:range()
-    node_indent = vim.fn.indent(srow+1)
+  local node = get_node_at_line(parser:parse():root(), lnum-1)
+  local indent_queries = get_indents(vim.api.nvim_get_current_buf())
+  local indents = indent_queries.indents
+  local branches = indent_queries.branches
+  if not indents then return 0 end
 
-    parent = locals.containing_scope(parent:parent() or parent) 
-    parent_row = parent:range()
-    parent_indent = vim.fn.indent(parent_row+1)
+  while node and branches[tostring(node)] do
+    node = node:parent()
+  end
 
-    if node_indent == 0 and parent_indent == 0 then
-      return 0
+  local ind = 0
+  while node do
+    node = node:parent()
+    if indents[tostring(node)] then
+      ind = ind + vim.bo.tabstop
     end
   end
 
-  local tabstop = vim.bo.tabstop
-  return parent_indent + tabstop
+  return ind
 end
 
 local indent_funcs = {}
diff --git a/lua/nvim-treesitter/query.lua b/lua/nvim-treesitter/query.lua
index 3b3b171ad..4b2aec3bf 100644
--- a/lua/nvim-treesitter/query.lua
+++ b/lua/nvim-treesitter/query.lua
@@ -8,7 +8,7 @@ local M = {}
 
 local query_cache = caching.create_buffer_cache()
 
-M.built_in_query_groups = {'highlights', 'locals', 'textobjects', 'folds'}
+M.built_in_query_groups = {'highlights', 'locals', 'textobjects', 'folds', 'indents'}
 
 -- Creates a function that checks whether a given query exists
 -- for a specific language.
diff --git a/lua/nvim-treesitter/ts_utils.lua b/lua/nvim-treesitter/ts_utils.lua
index 6647cc968..b235c0139 100644
--- a/lua/nvim-treesitter/ts_utils.lua
+++ b/lua/nvim-treesitter/ts_utils.lua
@@ -201,31 +201,26 @@ end
 
 --- Memoizes a function based on the buffer tick of the provided bufnr.
 -- The cache entry is cleared when the buffer is detached to avoid memory leaks.
--- @param fn: the fn to memoize
--- @param bufnr_fn: a function that receives all arguments passed to the function
---                  and returns the bufnr from the arguments
+-- @param fn: the fn to memoize, taking the bufnr as first argument
 -- @returns a memoized function
-function M.memoize_by_buf_tick(fn, bufnr_fn)
-  local bufnr_fn = bufnr_fn or function(a) return a end
+function M.memoize_by_buf_tick(fn)
   local cache = {}
 
-  return function(...)
-    local bufnr = bufnr_fn(...)
-    local tick = api.nvim_buf_get_changedtick(bufnr)
-
+  return function(bufnr)
     if cache[bufnr] then
-      if cache[bufnr].last_tick == tick then
-        return cache[bufnr].result
-      end
+      return cache[bufnr]
     else
       cache[bufnr] = {}
-      api.nvim_buf_attach(bufnr, false, { on_detach = function() cache[bufnr] = nil end })
+      api.nvim_buf_attach(bufnr, false,
+        {
+          on_changedtick = function() cache[bufnr] = fn(bufnr) end,
+          on_detach = function() cache[bufnr] = nil end
+        }
+      )
     end
 
-    cache[bufnr].last_tick = tick
-    cache[bufnr].result = fn(...)
-
-    return cache[bufnr].result
+    cache[bufnr] = fn(bufnr)
+    return cache[bufnr]
   end
 end
 
diff --git a/queries/lua/indents.scm b/queries/lua/indents.scm
new file mode 100644
index 000000000..e59b68085
--- /dev/null
+++ b/queries/lua/indents.scm
@@ -0,0 +1,23 @@
+[
+  (function_definition)
+  (local_function)
+  (function)
+  (if_statement)
+  (for_statement)
+  (for_in_statement)
+  (repeat_statement)
+  (while_statement)
+  (table)
+  (arguments)
+] @indent
+
+[
+  "end"
+  "{"
+  "}"
+  "("
+  ")"
+  "then"
+  (else)
+  (elseif)
+] @branch