feat: improve indent module

get_node_at_line should return appropriate child if available

5 files changed, 38 insertions, 11 deletions

diff --git a/lua/nvim-treesitter/indent.lua b/lua/nvim-treesitter/indent.lua
index ebf7695de..b585f5862 100644
--- a/lua/nvim-treesitter/indent.lua
+++ b/lua/nvim-treesitter/indent.lua
@@ -4,11 +4,35 @@ 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, _, erow = node:range()
+    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
 
@@ -89,7 +113,7 @@ function M.get_indent(lnum)
         -- 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 <= end_row then
+        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)
diff --git a/queries/lua/indents.scm b/queries/lua/indents.scm
index d0a16d20f..8203a6a49 100644
--- a/queries/lua/indents.scm
+++ b/queries/lua/indents.scm
@@ -13,18 +13,21 @@
   (arguments)
 ] @indent
 
- @ignore
-
 [
+  "do"
   "end"
+  "then"
   "until"
   "{"
   "}"
   "("
   ")"
-  "then"
-  (else_statement)
+  "elseif"
   (elseif_statement)
+  "else"
+  (else_statement)
 ] @branch
 
 (comment) @ignore
+
+(string) @ignore
diff --git a/tests/indent/lua_spec.lua b/tests/indent/lua_spec.lua
index 3b3beafea..cf615713c 100644
--- a/tests/indent/lua_spec.lua
+++ b/tests/indent/lua_spec.lua
@@ -28,7 +28,7 @@ describe("indent Lua:", function()
     run:new_line("table.lua", { on_line = 1, text = "b = 0,", indent = 2 })
     run:new_line("table.lua", { on_line = 5, text = "4,", indent = 4 })
     run:new_line("table.lua", { on_line = 7, text = "4,", indent = 4 })
-    run:new_line("loop.lua", { on_line = 4, text = "x = x + 1", indent = 2 }, "expected failure", XFAIL)
+    run:new_line("loop.lua", { on_line = 4, text = "x = x + 1", indent = 2 })
     run:new_line("cond.lua", { on_line = 5, text = "x = x + 1", indent = 2 })
     run:new_line("cond.lua", { on_line = 7, text = "x = x + 1", indent = 2 })
     run:new_line("cond.lua", { on_line = 8, text = "x = x + 1", indent = 4 })
diff --git a/tests/indent/python_spec.lua b/tests/indent/python_spec.lua
index f1769c4d8..18a1413a2 100644
--- a/tests/indent/python_spec.lua
+++ b/tests/indent/python_spec.lua
@@ -26,8 +26,8 @@ describe("indent Python:", function()
     run:new_line("basic_blocks.py", { on_line = 1, text = "wait,", indent = 4 })
     run:new_line("basic_blocks.py", { on_line = 6, text = "x += 1", indent = 4 })
     run:new_line("basic_blocks.py", { on_line = 10, text = "x += 1", indent = 8 })
-    run:new_line("basic_blocks.py", { on_line = 7, text = "x += 1", indent = 4 }, "7, after last line of a block")
-    run:new_line("basic_blocks.py", { on_line = 11, text = "x += 1", indent = 8 }, "11, after last line of a block")
+    run:new_line("basic_blocks.py", { on_line = 7, text = "x += 1", indent = 0 })
+    run:new_line("basic_blocks.py", { on_line = 11, text = "x += 1", indent = 4 })
     run:new_line("basic_collections.py", { on_line = 3, text = "4,", indent = 4 })
     run:new_line("comprehensions.py", { on_line = 8, text = "if x != 2", indent = 4 })
     run:new_line("control_flow.py", { on_line = 23, text = "x = 4", indent = 4 }, "expected failure", XFAIL)
diff --git a/tests/indent/rust/macro.rs b/tests/indent/rust/macro.rs
index 608e157fc..0900f4c12 100644
--- a/tests/indent/rust/macro.rs
+++ b/tests/indent/rust/macro.rs
@@ -2,10 +2,10 @@ macro_rules! foo {
     ($a:ident, $b:ident, $c:ident) => {
         struct $a;
         struct $b;
-    },
+    };
     ($a:ident) => {
         struct $a;
-    },
+    };
 }
 
 foo! {