feat: rewrite indent module

CONTRIBUTING.md

@@ -269,10 +269,11 @@ the node describing the language and `@content` to describe the injection region
 ### Indents
 
 ```
-@indent ; Indent when matching this node
-@branch ; Dedent when matching this node
-@return ; Dedent when matching this node
-@ignore ; Skip this node when calculating the indentation level
+@indent ; Indent children when matching this node
+@dedent ; Dedent children when matching this node
+@branch ; Dedent itself when matching this node
+@ignore ; Do not indent in this node
+@auto   ; Behaves like 'autoindent' buffer option
 ```
 
 [Zulip]: nvim-treesitter.zulipchat.com

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()
-
-  -- 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_()
-
-        -- 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
-    end
+  local indent = 0
+  if root:start() ~= 0 then
+    -- injected tree
+    indent = vim.fn.indent(root:start() + 1)
   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
-    end
-  end
-
-  while node and q.branches[node_fmt(node)] do
-    node = node:parent()
-  end
-
-  local first = true
-  local prev_row = node:start()
+  -- tracks to ensure multiple indent levels are not applied for same line
+  local is_processed_by_row = {}
 
   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
+    -- 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
 
-    -- 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
-      indent = indent + indent_size
-      prev_row = row
+    -- 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
 
+    local srow, _, erow = node:range()
+
+    local is_processed = false
+
+    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 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
+      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

queries/c/indents.scm

@@ -8,18 +8,38 @@
   (enumerator_list)
   (struct_specifier)
   (compound_literal_expression)
+  (parameter_list)
+  (initializer_list)
+  (concatenated_string)
+  (while_statement)
+  (for_statement)
+  (switch_statement)
 ] @indent
-
+(if_statement condition: (_) @indent)
+((if_statement
+  consequence: (_) @_consequence
+  (#not-has-type? @_consequence compound_statement)
+  ) @indent)
+(init_declarator 
+  value: [
+    (binary_expression)
+  ] @indent)
 
 [
   "#define"
   "#ifdef"
+  "#if"
+  "#else"
+  "else"
   "#endif"
-  "{"
+  ")"
   "}"
 ] @branch
 
 [
   (comment)
-  (preproc_function_def)
+  (preproc_arg)
+  (string_literal)
 ] @ignore
+
+(binary_expression) @auto

queries/cpp/indents.scm

@@ -1,7 +1,8 @@
 ; inherits: c
 
 [
-  (field_declaration_list)
   (class_specifier)
   (condition_clause)
 ] @indent
+
+(access_specifier) @branch

queries/lua/indents.scm

@@ -1,7 +1,6 @@
 [
   (function_definition)
   (function_declaration)
-  (variable_declaration)
   (field)
   (do_statement)
   (while_statement)
@@ -11,8 +10,13 @@
   (return_statement)
   (table_constructor)
   (arguments)
+  (return_statement)
 ] @indent
 
+(return_statement
+  (expression_list
+    (function_call))) @dedent
+
 [
   "do"
   "end"
@@ -30,4 +34,4 @@
 
 (comment) @ignore
 
-(string) @ignore
+(string) @auto

queries/python/indents.scm

@@ -23,8 +23,11 @@
   (parameters)
   (binary_operator)
 
+  (lambda)
   (function_definition)
   (class_definition)
+
+  (concatenated_string)
 ] @indent
 
 [
@@ -37,12 +40,4 @@
   (finally_clause)
 ] @branch
 
-[
-  (return_statement)
-  (pass_statement)
-  (raise_statement)
-] @return
-
-[
-  (string)
-] @ignore
+(string) @auto

queries/rust/indents.scm

@@ -16,19 +16,31 @@
   (block)
   (where_clause)
   (use_list)
+  (array_expression)
+  (ordered_field_declaration_list)
+  (field_declaration_list)
+  (enum_variant_list)
+  (parameters)
+  (token_tree)
+  (macro_definition)
 ] @indent
+(trait_item body: (_) @indent)
+(string_literal (escape_sequence)) @indent
+
+(impl_item (where_clause) @dedent)
 
 [
   "where"
-  "("
   ")"
-  "["
   "]"
-  "{"
   "}"
 ] @branch
+(impl_item (declaration_list) @branch)
 
 [
   (line_comment)
-  (raw_string_literal)
+  (string_literal)
 ] @ignore
+
+
+(raw_string_literal) @auto

tests/indent/c_spec.lua

@@ -13,15 +13,9 @@ describe("indent C:", function()
     runner:whole_file(".", {
       expected_failures = {
         "./ternary.c",
-        "./string.c",
         "./preproc_func.c",
-        "./preproc_cond.c",
-        "./no_braces.c",
         "./label.c",
-        "./func.c",
-        "./expr.c",
         "./comment.c",
-        "./array.c",
       },
     })
   end)
@@ -36,14 +30,14 @@ describe("indent C:", function()
     runner:new_line("label.c", { on_line = 3, text = "normal:", indent = 0 }, "expected failure", XFAIL)
     runner:new_line("loop.c", { on_line = 3, text = "x++;", indent = 8 })
     runner:new_line("preproc_cond.c", { on_line = 5, text = "x++;", indent = 4 })
-    runner:new_line("preproc_func.c", { on_line = 3, text = "x++; \\", indent = 8 })
-    runner:new_line("string.c", { on_line = 1, text = "brave new \\", indent = 0 }, "expected failure", XFAIL)
+    runner:new_line("preproc_func.c", { on_line = 3, text = "x++; \\", indent = 8 }, "expected failure", XFAIL)
+    runner:new_line("string.c", { on_line = 1, text = "brave new \\", indent = 0 })
     runner:new_line("string.c", { on_line = 4, text = '"brave new "', indent = 4 })
     runner:new_line("struct.c", { on_line = 4, text = "int y;", indent = 8 })
     runner:new_line("switch.c", { on_line = 3, text = "x++;", indent = 12 })
-    runner:new_line("ternary.c", { on_line = 4, text = ": (x == 0) : 0", indent = 8 })
+    runner:new_line("ternary.c", { on_line = 4, text = ": (x == 0) : 0", indent = 8 }, "expected failure", XFAIL)
     -- the line after inserted one will be left with wrong indent but we only care about the inserted one
-    runner:new_line("no_braces.c", { on_line = 4, text = "x++;", indent = 8 }, "expected failure", XFAIL)
+    runner:new_line("no_braces.c", { on_line = 4, text = "x++;", indent = 8 })
     runner:new_line("no_braces.c", { on_line = 7, text = "x++;", indent = 8 })
     runner:new_line("no_braces.c", { on_line = 10, text = "x++;", indent = 8 })
   end)

tests/indent/cpp_spec.lua

@@ -16,18 +16,9 @@ describe("indent C++:", function()
       expected_failures = {
         -- C
         "c/ternary.c",
-        "c/string.c",
         "c/preproc_func.c",
-        "c/preproc_cond.c",
-        "c/no_braces.c",
         "c/label.c",
-        "c/func.c",
-        "c/expr.c",
         "c/comment.c",
-        "c/array.c",
-        -- C++
-        "cpp/access.cpp",
-        "cpp/stream.cpp",
       },
     })
   end)
@@ -46,14 +37,14 @@ describe("indent C++:", function()
     run:new_line("c/label.c", { on_line = 3, text = "normal:", indent = 0 }, "expected failure", XFAIL)
     run:new_line("c/loop.c", { on_line = 3, text = "x++;", indent = 8 })
     run:new_line("c/preproc_cond.c", { on_line = 5, text = "x++;", indent = 4 })
-    run:new_line("c/preproc_func.c", { on_line = 3, text = "x++; \\", indent = 8 })
-    run:new_line("c/string.c", { on_line = 1, text = "brave new \\", indent = 0 }, "expected failure", XFAIL)
+    run:new_line("c/preproc_func.c", { on_line = 3, text = "x++; \\", indent = 8 }, "expected failure", XFAIL)
+    run:new_line("c/string.c", { on_line = 1, text = "brave new \\", indent = 0 })
     run:new_line("c/string.c", { on_line = 4, text = '"brave new "', indent = 4 })
     run:new_line("c/struct.c", { on_line = 4, text = "int y;", indent = 8 })
     run:new_line("c/switch.c", { on_line = 3, text = "x++;", indent = 12 })
-    run:new_line("c/ternary.c", { on_line = 4, text = ": (x == 0) : 0", indent = 8 })
+    run:new_line("c/ternary.c", { on_line = 4, text = ": (x == 0) : 0", indent = 8 }, "expected failure", XFAIL)
     -- the line after inserted one will be left with wrong indent but we only care about the inserted one
-    run:new_line("c/no_braces.c", { on_line = 4, text = "x++;", indent = 8 }, "expected failure", XFAIL)
+    run:new_line("c/no_braces.c", { on_line = 4, text = "x++;", indent = 8 })
     run:new_line("c/no_braces.c", { on_line = 7, text = "x++;", indent = 8 })
     run:new_line("c/no_braces.c", { on_line = 10, text = "x++;", indent = 8 })
   end)

tests/indent/lua_spec.lua

@@ -10,21 +10,23 @@ local run = Runner:new(it, "tests/indent/lua", {
 
 describe("indent Lua:", function()
   describe("whole file:", function()
-    run:whole_file(".", { expected_failures = {
-      "./comment.lua",
-    } })
+    run:whole_file(".", {
+      expected_failures = {
+        "./comment.lua",
+      },
+    })
   end)
 
   describe("new line:", function()
     run:new_line("comment.lua", { on_line = 1, text = "line", indent = "-- " })
-    run:new_line("comment.lua", { on_line = 5, text = "multiline", indent = "  " })
+    run:new_line("comment.lua", { on_line = 5, text = "multiline", indent = "  " }, "expected failure", XFAIL)
     run:new_line("func.lua", { on_line = 1, text = "x = x + 1", indent = 2 })
     run:new_line("func.lua", { on_line = 2, text = "y = y + 1", indent = 4 })
     run:new_line("func.lua", { on_line = 5, text = "3,", indent = 4 })
-    run:new_line("string.lua", { on_line = 1, text = "x", indent = 0 }, "expected failure", XFAIL)
-    run:new_line("string.lua", { on_line = 2, text = "x", indent = 0 }, "expected failure", XFAIL)
+    run:new_line("string.lua", { on_line = 1, text = "x", indent = 0 })
+    run:new_line("string.lua", { on_line = 2, text = "x", indent = 0 })
     run:new_line("string.lua", { on_line = 3, text = "x", indent = 2 })
-    run:new_line("string.lua", { on_line = 4, text = "x", indent = 4 }, "expected failure", XFAIL)
+    run:new_line("string.lua", { on_line = 4, text = "x", indent = 4 })
     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 })

tests/indent/python_spec.lua

@@ -15,7 +15,6 @@ describe("indent Python:", function()
         "./aligned_indent.py",
         "./branches.py",
         "./hanging_indent.py",
-        "./join_lines.py",
         "./nested_collections.py",
       },
     })
@@ -26,8 +25,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 = 0 })
-    run:new_line("basic_blocks.py", { on_line = 11, text = "x += 1", indent = 4 })
+    run:new_line("basic_blocks.py", { on_line = 7, text = "x += 1", indent = 4 })
+    run:new_line("basic_blocks.py", { on_line = 11, text = "x += 1", indent = 8 })
     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)
@@ -39,9 +38,9 @@ describe("indent Python:", function()
     run:new_line("nested_collections.py", { on_line = 5, text = "0,", indent = 12 })
     run:new_line("nested_collections.py", { on_line = 6, text = ",0", indent = 12 })
     run:new_line("nested_collections.py", { on_line = 29, text = "[1, 2],", indent = 12 })
-    run:new_line("nested_collections.py", { on_line = 39, text = "0,", indent = 5 })
+    run:new_line("nested_collections.py", { on_line = 39, text = "0,", indent = 5 }, "expected failure", XFAIL)
     run:new_line("strings.py", { on_line = 14, text = "x", indent = 4 })
-    run:new_line("strings.py", { on_line = 15, text = "x", indent = 0 }, nil, XFAIL)
-    run:new_line("strings.py", { on_line = 16, text = "x", indent = 8 }, nil, XFAIL)
+    run:new_line("strings.py", { on_line = 15, text = "x", indent = 0 })
+    run:new_line("strings.py", { on_line = 16, text = "x", indent = 8 })
   end)
 end)

tests/indent/rust/macro.rs

@@ -1,10 +1,10 @@
 macro_rules! foo {
     ($a:ident, $b:ident, $c:ident) => {
-        struct $a;
-        struct $b;
+        struct a { value: $a };
+        struct b { value: $b };
     };
     ($a:ident) => {
-        struct $a;
+        struct a { value: $a };
     };
 }
 

tests/indent/rust_spec.lua

@@ -1,5 +1,4 @@
 local Runner = require("tests.indent.common").Runner
-local XFAIL = require("tests.indent.common").XFAIL
 
 local run = Runner:new(it, "tests/indent/rust", {
   tabstop = 4,
@@ -10,17 +9,7 @@ local run = Runner:new(it, "tests/indent/rust", {
 
 describe("indent Rust:", function()
   describe("whole file:", function()
-    run:whole_file(".", {
-      expected_failures = {
-        "./enum.rs",
-        "./func.rs",
-        "./array.rs",
-        "./where.rs",
-        "./trait.rs",
-        "./string.rs",
-        "./macro.rs",
-      },
-    })
+    run:whole_file "."
   end)
 
   describe("new line:", function()
@@ -29,8 +18,8 @@ describe("indent Rust:", function()
     run:new_line("comment.rs", { on_line = 3, text = "a", indent = "/// " })
     run:new_line("cond.rs", { on_line = 11, text = "x += 1;", indent = 12 })
     run:new_line("cond.rs", { on_line = 2, text = "x += 1;", indent = 8 })
-    run:new_line("cond.rs", { on_line = 4, text = "x += 1;", indent = 8 }, "expected_failures", XFAIL)
-    run:new_line("cond.rs", { on_line = 6, text = "x += 1;", indent = 8 }, "expected_failures", XFAIL)
+    run:new_line("cond.rs", { on_line = 4, text = "x += 1;", indent = 8 })
+    run:new_line("cond.rs", { on_line = 6, text = "x += 1;", indent = 8 })
     run:new_line("enum.rs", { on_line = 2, text = "Q,", indent = 4 })
     run:new_line("enum.rs", { on_line = 4, text = "i32,", indent = 8 })
     run:new_line("enum.rs", { on_line = 8, text = "z: u32,", indent = 8 })
@@ -49,7 +38,7 @@ describe("indent Rust:", function()
     run:new_line("mod.rs", { on_line = 1, text = "const Z: i32 = 1;", indent = 4 })
     run:new_line("mod.rs", { on_line = 2, text = "const Z: i32 = 1;", indent = 4 })
     run:new_line("mod.rs", { on_line = 6, text = "const Z: i32 = 1;", indent = 8 })
-    run:new_line("string.rs", { on_line = 2, text = "brave new", indent = 0 }, "expected_failures", XFAIL)
+    run:new_line("string.rs", { on_line = 2, text = "brave new", indent = 0 })
     run:new_line("string.rs", { on_line = 5, text = "brave new \\", indent = 8 })
     run:new_line("string.rs", { on_line = 9, text = "brave new \\", indent = 8 })
     run:new_line("struct.rs", { on_line = 1, text = "z: i32,", indent = 4 })