feat: rewrite indent module

1 files changed, 47 insertions, 101 deletions

diff --git a/lua/nvim-treesitter/indent.lua b/lua/nvim-treesitter/indent.lua
index b585f5862..7f7891261 100644
--- a/lua/nvim-treesitter/indent.lua
+++ b/lua/nvim-treesitter/indent.lua
@@ -4,74 +4,33 @@ local tsutils = require "nvim-treesitter.ts_utils"
 
 local M = {}
 
----@param lnum number (0-indexed)
-local function get_last_node_at_line(root, lnum)
-  local node
-  for i = 0, root:child_count() - 1 do
-    local child = root:child(i)
-    local child_srow = child:start()
-    if child_srow > lnum then
-      break
-    end
-    if child_srow == lnum then
-      node = child
-    end
-  end
-  return node
-end
-
--- TODO(kiyan): move this in tsutils and document it
----@param lnum number (0-indexed)
-local function get_node_at_line(root, lnum)
-  for node in root:iter_children() do
-    local srow, scol, erow = node:range()
-    if srow == lnum then
-      if node:child_count() > 0 then
-        local child = get_last_node_at_line(node, srow)
-        if child and child:named() and ({ child:start() })[2] == scol then
-          -- last child node is named and start at the same col as parent
-          return child
-        end
-      end
-      return node
-    end
-
-    if node:child_count() > 0 and srow < lnum and lnum <= erow then
-      return get_node_at_line(node, lnum)
-    end
-  end
-end
-
-local function node_fmt(node)
-  if not node then
-    return nil
-  end
-  return tostring(node)
-end
-
 local get_indents = tsutils.memoize_by_buf_tick(function(bufnr, root, lang)
   local get_map = function(capture)
     local matches = queries.get_capture_matches(bufnr, capture, "indents", root, lang) or {}
     local map = {}
     for _, node in ipairs(matches) do
-      map[tostring(node)] = true
+      map[node:id()] = true
     end
     return map
   end
 
   return {
+    autos = get_map "@auto.node",
     indents = get_map "@indent.node",
+    dedents = get_map "@dedent.node",
     branches = get_map "@branch.node",
-    returns = get_map "@return.node",
     ignores = get_map "@ignore.node",
+    aligned_indents = get_map "@aligned_indent.node",
+    hanging_indents = get_map "@hanging_indent.node",
   }
 end, {
   -- Memoize by bufnr and lang together.
-  key = function(bufnr, _, lang)
-    return tostring(bufnr) .. "_" .. lang
+  key = function(bufnr, root, lang)
+    return tostring(bufnr) .. root:id() .. "_" .. lang
   end,
 })
 
+---@param lnum number (1-indexed)
 function M.get_indent(lnum)
   local parser = parsers.get_parser()
   if not parser or not lnum then
@@ -87,74 +46,61 @@ function M.get_indent(lnum)
   end
 
   local q = get_indents(vim.api.nvim_get_current_buf(), root, lang_tree:lang())
-  local node = get_node_at_line(root, lnum - 1)
+  local is_empty_line = string.match(vim.fn.getline(lnum), "^%s*$") ~= nil
+  local node
+  if is_empty_line then
+    local prevlnum = vim.fn.prevnonblank(lnum)
+    local col = #vim.fn.getline(prevlnum) - 1
+    node = root:descendant_for_range(prevlnum - 1, col, prevlnum - 1, col)
+  else
+    local col = vim.fn.indent(lnum)
+    node = root:descendant_for_range(lnum - 1, col, lnum - 1, col)
+  end
 
-  local indent = 0
   local indent_size = vim.fn.shiftwidth()
+  local indent = 0
+  if root:start() ~= 0 then
+    -- injected tree
+    indent = vim.fn.indent(root:start() + 1)
+  end
 
-  -- to get correct indentation when we land on an empty line (for instance by typing `o`), we try
-  -- to use indentation of previous nonblank line, this solves the issue also for languages that
-  -- do not use @branch after blocks (e.g. Python)
-  if not node then
-    local prevnonblank = vim.fn.prevnonblank(lnum)
-    if prevnonblank ~= lnum then
-      local prev_node = get_node_at_line(root, prevnonblank - 1)
-      -- get previous node in any case to avoid erroring
-      while not prev_node and prevnonblank - 1 > 0 do
-        prevnonblank = vim.fn.prevnonblank(prevnonblank - 1)
-        prev_node = get_node_at_line(root, prevnonblank - 1)
-      end
-
-      -- nodes can be marked @return to prevent using them
-      if prev_node and not q.returns[node_fmt(prev_node)] then
-        local row = prev_node:start()
-        local end_row = prev_node:end_()
+  -- tracks to ensure multiple indent levels are not applied for same line
+  local is_processed_by_row = {}
 
-        -- if the previous node is being constructed (like function() `o` in lua), or line is inside the node
-        -- we indent one more from the start of node, else we indent default
-        -- NOTE: this doesn't work for python which behave strangely
-        if prev_node:has_error() or lnum - 1 < end_row then
-          return vim.fn.indent(row + 1) + indent_size
-        end
-        return vim.fn.indent(row + 1)
-      end
+  while node do
+    -- do 'autoindent' if not marked as @indent
+    if not q.indents[node:id()] and q.autos[node:id()] and node:start() < lnum - 1 and lnum - 1 <= node:end_() then
+      return -1
     end
-  end
 
-  -- if the prevnonblank fails (prev_node wraps our line) we need to fall back to taking
-  -- the first child of the node that wraps the current line, or the wrapper itself
-  if not node then
-    local wrapper = root:descendant_for_range(lnum - 1, 0, lnum - 1, -1)
-    node = wrapper:child(0) or wrapper
-    if q.indents[node_fmt(wrapper)] ~= nil and wrapper ~= root then
-      indent = indent_size
+    -- Do not indent if we are inside an @ignore block.
+    -- If a node spans from L1,C1 to L2,C2, we know that lines where L1 < line <= L2 would
+    -- have their indentations contained by the node.
+    if not q.indents[node:id()] and q.ignores[node:id()] and node:start() < lnum - 1 and lnum - 1 <= node:end_() then
+      return 0
     end
-  end
 
-  while node and q.branches[node_fmt(node)] do
-    node = node:parent()
-  end
+    local srow, _, erow = node:range()
 
-  local first = true
-  local prev_row = node:start()
+    local is_processed = false
 
-  while node do
-    -- Do not indent if we are inside an @ignore block.
-    -- If a node spans from L1,C1 to L2,C2, we know that lines where L1 < line <= L2 would
-    -- have their indentations contained by the node.
-    if q.ignores[node_fmt(node)] and node:start() < lnum - 1 and lnum - 1 <= node:end_() then
-      return -1
+    if
+      not is_processed_by_row[srow]
+      and ((q.branches[node:id()] and srow == lnum - 1) or (q.dedents[node:id()] and srow ~= lnum - 1))
+    then
+      indent = indent - indent_size
+      is_processed = true
     end
 
-    -- do not indent the starting node, do not add multiple indent levels on single line
-    local row = node:start()
-    if not first and q.indents[node_fmt(node)] and prev_row ~= row then
+    -- do not indent for nodes that starts-and-ends on same line and starts on target line (lnum)
+    if not is_processed_by_row[srow] and (q.indents[node:id()] and srow ~= erow and srow ~= lnum - 1) then
       indent = indent + indent_size
-      prev_row = row
+      is_processed = true
     end
 
+    is_processed_by_row[srow] = is_processed_by_row[srow] or is_processed
+
     node = node:parent()
-    first = false
   end
 
   return indent