Merge remote-tracking branch 'upstream/master' into userreg

9 files changed, 861 insertions, 674 deletions

diff --git a/runtime/lua/vim/treesitter/_fold.lua b/runtime/lua/vim/treesitter/_fold.lua
index 5c1cc06908..d96cc966de 100644
--- a/runtime/lua/vim/treesitter/_fold.lua
+++ b/runtime/lua/vim/treesitter/_fold.lua
@@ -5,35 +5,20 @@ local Range = require('vim.treesitter._range')
 local api = vim.api
 
 ---@class TS.FoldInfo
----@field levels table<integer,string>
----@field levels0 table<integer,integer>
----@field private start_counts table<integer,integer>
----@field private stop_counts table<integer,integer>
+---@field levels string[] the foldexpr result for each line
+---@field levels0 integer[] the raw fold levels
+---@field edits? {[1]: integer, [2]: integer} line range edited since the last invocation of the callback scheduled in on_bytes. 0-indexed, end-exclusive.
 local FoldInfo = {}
 FoldInfo.__index = FoldInfo
 
 ---@private
 function FoldInfo.new()
   return setmetatable({
-    start_counts = {},
-    stop_counts = {},
     levels0 = {},
     levels = {},
   }, FoldInfo)
 end
 
----@package
----@param srow integer
----@param erow integer
-function FoldInfo:invalidate_range(srow, erow)
-  for i = srow, erow do
-    self.start_counts[i + 1] = nil
-    self.stop_counts[i + 1] = nil
-    self.levels0[i + 1] = nil
-    self.levels[i + 1] = nil
-  end
-end
-
 --- Efficiently remove items from middle of a list a list.
 ---
 --- Calling table.remove() in a loop will re-index the tail of the table on
@@ -55,12 +40,10 @@ end
 
 ---@package
 ---@param srow integer
----@param erow integer
+---@param erow integer 0-indexed, exclusive
 function FoldInfo:remove_range(srow, erow)
   list_remove(self.levels, srow + 1, erow)
   list_remove(self.levels0, srow + 1, erow)
-  list_remove(self.start_counts, srow + 1, erow)
-  list_remove(self.stop_counts, srow + 1, erow)
 end
 
 --- Efficiently insert items into the middle of a list.
@@ -91,46 +74,37 @@ end
 
 ---@package
 ---@param srow integer
----@param erow integer
+---@param erow integer 0-indexed, exclusive
 function FoldInfo:add_range(srow, erow)
-  list_insert(self.levels, srow + 1, erow, '-1')
+  list_insert(self.levels, srow + 1, erow, '=')
   list_insert(self.levels0, srow + 1, erow, -1)
-  list_insert(self.start_counts, srow + 1, erow, nil)
-  list_insert(self.stop_counts, srow + 1, erow, nil)
 end
 
 ---@package
----@param lnum integer
-function FoldInfo:add_start(lnum)
-  self.start_counts[lnum] = (self.start_counts[lnum] or 0) + 1
-end
-
----@package
----@param lnum integer
-function FoldInfo:add_stop(lnum)
-  self.stop_counts[lnum] = (self.stop_counts[lnum] or 0) + 1
-end
-
----@package
----@param lnum integer
----@return integer
-function FoldInfo:get_start(lnum)
-  return self.start_counts[lnum] or 0
+---@param srow integer
+---@param erow_old integer
+---@param erow_new integer 0-indexed, exclusive
+function FoldInfo:edit_range(srow, erow_old, erow_new)
+  if self.edits then
+    self.edits[1] = math.min(srow, self.edits[1])
+    if erow_old <= self.edits[2] then
+      self.edits[2] = self.edits[2] + (erow_new - erow_old)
+    end
+    self.edits[2] = math.max(self.edits[2], erow_new)
+  else
+    self.edits = { srow, erow_new }
+  end
 end
 
 ---@package
----@param lnum integer
----@return integer
-function FoldInfo:get_stop(lnum)
-  return self.stop_counts[lnum] or 0
-end
-
-local function trim_level(level)
-  local max_fold_level = vim.wo.foldnestmax
-  if level > max_fold_level then
-    return max_fold_level
+---@return integer? srow
+---@return integer? erow 0-indexed, exclusive
+function FoldInfo:flush_edit()
+  if self.edits then
+    local srow, erow = self.edits[1], self.edits[2]
+    self.edits = nil
+    return srow, erow
   end
-  return level
 end
 
 --- If a parser doesn't have any ranges explicitly set, treesitter will
@@ -140,10 +114,10 @@ end
 --- TODO(lewis6991): Handle this generally
 ---
 --- @param bufnr integer
---- @param erow integer?
+--- @param erow integer? 0-indexed, exclusive
 --- @return integer
 local function normalise_erow(bufnr, erow)
-  local max_erow = api.nvim_buf_line_count(bufnr) - 1
+  local max_erow = api.nvim_buf_line_count(bufnr)
   return math.min(erow or max_erow, max_erow)
 end
 
@@ -152,31 +126,30 @@ end
 ---@param bufnr integer
 ---@param info TS.FoldInfo
 ---@param srow integer?
----@param erow integer?
+---@param erow integer? 0-indexed, exclusive
 ---@param parse_injections? boolean
 local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
   srow = srow or 0
   erow = normalise_erow(bufnr, erow)
 
-  info:invalidate_range(srow, erow)
-
-  local prev_start = -1
-  local prev_stop = -1
-
   local parser = ts.get_parser(bufnr)
 
   parser:parse(parse_injections and { srow, erow } or nil)
 
+  local enter_counts = {} ---@type table<integer, integer>
+  local leave_counts = {} ---@type table<integer, integer>
+  local prev_start = -1
+  local prev_stop = -1
+
   parser:for_each_tree(function(tree, ltree)
     local query = ts.query.get(ltree:lang(), 'folds')
     if not query then
       return
     end
 
-    -- erow in query is end-exclusive
-    local q_erow = erow and erow + 1 or -1
-
-    for id, node, metadata in query:iter_captures(tree:root(), bufnr, srow, q_erow) do
+    -- Collect folds starting from srow - 1, because we should first subtract the folds that end at
+    -- srow - 1 from the level of srow - 1 to get accurate level of srow.
+    for id, node, metadata in query:iter_captures(tree:root(), bufnr, math.max(srow - 1, 0), erow) do
       if query.captures[id] == 'fold' then
         local range = ts.get_range(node, bufnr, metadata[id])
         local start, _, stop, stop_col = Range.unpack4(range)
@@ -193,8 +166,8 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
         if
           fold_length > vim.wo.foldminlines and not (start == prev_start and stop == prev_stop)
         then
-          info:add_start(start + 1)
-          info:add_stop(stop + 1)
+          enter_counts[start + 1] = (enter_counts[start + 1] or 0) + 1
+          leave_counts[stop + 1] = (leave_counts[stop + 1] or 0) + 1
           prev_start = start
           prev_stop = stop
         end
@@ -202,16 +175,15 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
     end
   end)
 
-  local current_level = info.levels0[srow] or 0
+  local nestmax = vim.wo.foldnestmax
+  local level0_prev = info.levels0[srow] or 0
+  local leave_prev = leave_counts[srow] or 0
 
   -- We now have the list of fold opening and closing, fill the gaps and mark where fold start
-  for lnum = srow + 1, erow + 1 do
-    local last_trimmed_level = trim_level(current_level)
-    current_level = current_level + info:get_start(lnum)
-    info.levels0[lnum] = current_level
-
-    local trimmed_level = trim_level(current_level)
-    current_level = current_level - info:get_stop(lnum)
+  for lnum = srow + 1, erow do
+    local enter_line = enter_counts[lnum] or 0
+    local leave_line = leave_counts[lnum] or 0
+    local level0 = level0_prev - leave_prev + enter_line
 
     -- Determine if it's the start/end of a fold
     -- NB: vim's fold-expr interface does not have a mechanism to indicate that
@@ -219,14 +191,36 @@ local function get_folds_levels(bufnr, info, srow, erow, parse_injections)
     --  ( \n ( \n )) \n (( \n ) \n )
     -- versus
     --  ( \n ( \n ) \n ( \n ) \n )
-    -- If it did have such a mechanism, (trimmed_level - last_trimmed_level)
+    -- Both are represented by ['>1', '>2', '2', '>2', '2', '1'], and
+    -- vim interprets as the second case.
+    -- If it did have such a mechanism, (clamped - clamped_prev)
     -- would be the correct number of starts to pass on.
+    local adjusted = level0 ---@type integer
     local prefix = ''
-    if trimmed_level - last_trimmed_level > 0 then
+    if enter_line > 0 then
       prefix = '>'
+      if leave_line > 0 then
+        -- If this line ends a fold f1 and starts a fold f2, then move f1's end to the previous line
+        -- so that f2 gets the correct level on this line. This may reduce the size of f1 below
+        -- foldminlines, but we don't handle it for simplicity.
+        adjusted = level0 - leave_line
+        leave_line = 0
+      end
+    end
+
+    -- Clamp at foldnestmax.
+    local clamped = adjusted
+    if adjusted > nestmax then
+      prefix = ''
+      clamped = nestmax
     end
 
-    info.levels[lnum] = prefix .. tostring(trimmed_level)
+    -- Record the "real" level, so that it can be used as "base" of later get_folds_levels().
+    info.levels0[lnum] = adjusted
+    info.levels[lnum] = prefix .. tostring(clamped)
+
+    leave_prev = leave_line
+    level0_prev = adjusted
   end
 end
 
@@ -296,8 +290,12 @@ end
 local function on_changedtree(bufnr, foldinfo, tree_changes)
   schedule_if_loaded(bufnr, function()
     for _, change in ipairs(tree_changes) do
-      local srow, _, erow = Range.unpack4(change)
-      get_folds_levels(bufnr, foldinfo, srow, erow)
+      local srow, _, erow, ecol = Range.unpack4(change)
+      if ecol > 0 then
+        erow = erow + 1
+      end
+      -- Start from `srow - foldminlines`, because this edit may have shrunken the fold below limit.
+      get_folds_levels(bufnr, foldinfo, math.max(srow - vim.wo.foldminlines, 0), erow)
     end
     if #tree_changes > 0 then
       foldupdate(bufnr)
@@ -309,19 +307,46 @@ end
 ---@param foldinfo TS.FoldInfo
 ---@param start_row integer
 ---@param old_row integer
+---@param old_col integer
 ---@param new_row integer
-local function on_bytes(bufnr, foldinfo, start_row, old_row, new_row)
-  local end_row_old = start_row + old_row
-  local end_row_new = start_row + new_row
+---@param new_col integer
+local function on_bytes(bufnr, foldinfo, start_row, start_col, old_row, old_col, new_row, new_col)
+  -- extend the end to fully include the range
+  local end_row_old = start_row + old_row + 1
+  local end_row_new = start_row + new_row + 1
 
   if new_row ~= old_row then
+    -- foldexpr can be evaluated before the scheduled callback is invoked. So it may observe the
+    -- outdated levels, which may spuriously open the folds that didn't change. So we should shift
+    -- folds as accurately as possible. For this to be perfectly accurate, we should track the
+    -- actual TSNodes that account for each fold, and compare the node's range with the edited
+    -- range. But for simplicity, we just check whether the start row is completely removed (e.g.,
+    -- `dd`) or shifted (e.g., `o`).
     if new_row < old_row then
-      foldinfo:remove_range(end_row_new, end_row_old)
+      if start_col == 0 and new_row == 0 and new_col == 0 then
+        foldinfo:remove_range(start_row, start_row + (end_row_old - end_row_new))
+      else
+        foldinfo:remove_range(end_row_new, end_row_old)
+      end
     else
-      foldinfo:add_range(start_row, end_row_new)
+      if start_col == 0 and old_row == 0 and old_col == 0 then
+        foldinfo:add_range(start_row, start_row + (end_row_new - end_row_old))
+      else
+        foldinfo:add_range(end_row_old, end_row_new)
+      end
     end
+    foldinfo:edit_range(start_row, end_row_old, end_row_new)
+
+    -- This callback must not use on_bytes arguments, because they can be outdated when the callback
+    -- is invoked. For example, `J` with non-zero count triggers multiple on_bytes before executing
+    -- the scheduled callback. So we should collect the edits.
     schedule_if_loaded(bufnr, function()
-      get_folds_levels(bufnr, foldinfo, start_row, end_row_new)
+      local srow, erow = foldinfo:flush_edit()
+      if not srow then
+        return
+      end
+      -- Start from `srow - foldminlines`, because this edit may have shrunken the fold below limit.
+      get_folds_levels(bufnr, foldinfo, math.max(srow - vim.wo.foldminlines, 0), erow)
       foldupdate(bufnr)
     end)
   end
@@ -348,8 +373,8 @@ function M.foldexpr(lnum)
         on_changedtree(bufnr, foldinfos[bufnr], tree_changes)
       end,
 
-      on_bytes = function(_, _, start_row, _, _, old_row, _, _, new_row, _, _)
-        on_bytes(bufnr, foldinfos[bufnr], start_row, old_row, new_row)
+      on_bytes = function(_, _, start_row, start_col, _, old_row, old_col, _, new_row, new_col, _)
+        on_bytes(bufnr, foldinfos[bufnr], start_row, start_col, old_row, old_col, new_row, new_col)
       end,
 
       on_detach = function()
@@ -361,96 +386,15 @@ function M.foldexpr(lnum)
   return foldinfos[bufnr].levels[lnum] or '0'
 end
 
----@package
----@return { [1]: string, [2]: string[] }[]|string
-function M.foldtext()
-  local foldstart = vim.v.foldstart
-  local bufnr = api.nvim_get_current_buf()
-
-  ---@type boolean, LanguageTree
-  local ok, parser = pcall(ts.get_parser, bufnr)
-  if not ok then
-    return vim.fn.foldtext()
-  end
-
-  local query = ts.query.get(parser:lang(), 'highlights')
-  if not query then
-    return vim.fn.foldtext()
-  end
-
-  local tree = parser:parse({ foldstart - 1, foldstart })[1]
-
-  local line = api.nvim_buf_get_lines(bufnr, foldstart - 1, foldstart, false)[1]
-  if not line then
-    return vim.fn.foldtext()
-  end
-
-  ---@type { [1]: string, [2]: string[], range: { [1]: integer, [2]: integer } }[] | { [1]: string, [2]: string[] }[]
-  local result = {}
-
-  local line_pos = 0
-
-  for id, node, metadata in query:iter_captures(tree:root(), 0, foldstart - 1, foldstart) do
-    local name = query.captures[id]
-    local start_row, start_col, end_row, end_col = node:range()
-
-    local priority = tonumber(metadata.priority or vim.highlight.priorities.treesitter)
-
-    if start_row == foldstart - 1 and end_row == foldstart - 1 then
-      -- check for characters ignored by treesitter
-      if start_col > line_pos then
-        table.insert(result, {
-          line:sub(line_pos + 1, start_col),
-          {},
-          range = { line_pos, start_col },
-        })
-      end
-      line_pos = end_col
-
-      local text = line:sub(start_col + 1, end_col)
-      table.insert(result, { text, { { '@' .. name, priority } }, range = { start_col, end_col } })
-    end
-  end
-
-  local i = 1
-  while i <= #result do
-    -- find first capture that is not in current range and apply highlights on the way
-    local j = i + 1
-    while
-      j <= #result
-      and result[j].range[1] >= result[i].range[1]
-      and result[j].range[2] <= result[i].range[2]
-    do
-      for k, v in ipairs(result[i][2]) do
-        if not vim.tbl_contains(result[j][2], v) then
-          table.insert(result[j][2], k, v)
-        end
-      end
-      j = j + 1
-    end
-
-    -- remove the parent capture if it is split into children
-    if j > i + 1 then
-      table.remove(result, i)
-    else
-      -- highlights need to be sorted by priority, on equal prio, the deeper nested capture (earlier
-      -- in list) should be considered higher prio
-      if #result[i][2] > 1 then
-        table.sort(result[i][2], function(a, b)
-          return a[2] < b[2]
-        end)
-      end
-
-      result[i][2] = vim.tbl_map(function(tbl)
-        return tbl[1]
-      end, result[i][2])
-      result[i] = { result[i][1], result[i][2] }
-
-      i = i + 1
+api.nvim_create_autocmd('OptionSet', {
+  pattern = { 'foldminlines', 'foldnestmax' },
+  desc = 'Refresh treesitter folds',
+  callback = function()
+    for _, bufnr in ipairs(vim.tbl_keys(foldinfos)) do
+      foldinfos[bufnr] = FoldInfo.new()
+      get_folds_levels(bufnr, foldinfos[bufnr])
+      foldupdate(bufnr)
     end
-  end
-
-  return result
-end
-
+  end,
+})
 return M
diff --git a/runtime/lua/vim/treesitter/_meta.lua b/runtime/lua/vim/treesitter/_meta.lua
index 80c998b555..19d97d2820 100644
--- a/runtime/lua/vim/treesitter/_meta.lua
+++ b/runtime/lua/vim/treesitter/_meta.lua
@@ -1,4 +1,5 @@
 ---@meta
+error('Cannot require a meta file')
 
 ---@class TSNode: userdata
 ---@field id fun(self: TSNode): string
@@ -33,27 +34,26 @@
 ---@field byte_length fun(self: TSNode): integer
 local TSNode = {}
 
----@param query userdata
+---@param query TSQuery
 ---@param captures true
 ---@param start? integer
 ---@param end_? integer
 ---@param opts? table
----@return fun(): integer, TSNode, any
+---@return fun(): integer, TSNode, vim.treesitter.query.TSMatch
 function TSNode:_rawquery(query, captures, start, end_, opts) end
 
----@param query userdata
+---@param query TSQuery
 ---@param captures false
 ---@param start? integer
 ---@param end_? integer
 ---@param opts? table
----@return fun(): string, any
+---@return fun(): integer, vim.treesitter.query.TSMatch
 function TSNode:_rawquery(query, captures, start, end_, opts) end
 
 ---@alias TSLoggerCallback fun(logtype: 'parse'|'lex', msg: string)
 
----@class TSParser
----@field parse fun(self: TSParser, tree: TSTree?, source: integer|string, include_bytes: true): TSTree, Range6[]
----@field parse fun(self: TSParser, tree: TSTree?, source: integer|string, include_bytes: false|nil): TSTree, Range4[]
+---@class TSParser: userdata
+---@field parse fun(self: TSParser, tree: TSTree?, source: integer|string, include_bytes: boolean): TSTree, (Range4|Range6)[]
 ---@field reset fun(self: TSParser)
 ---@field included_ranges fun(self: TSParser, include_bytes: boolean?): integer[]
 ---@field set_included_ranges fun(self: TSParser, ranges: (Range6|TSNode)[])
@@ -62,19 +62,31 @@ function TSNode:_rawquery(query, captures, start, end_, opts) end
 ---@field _set_logger fun(self: TSParser, lex: boolean, parse: boolean, cb: TSLoggerCallback)
 ---@field _logger fun(self: TSParser): TSLoggerCallback
 
----@class TSTree
+---@class TSTree: userdata
 ---@field root fun(self: TSTree): TSNode
 ---@field edit fun(self: TSTree, _: integer, _: integer, _: integer, _: integer, _: integer, _: integer, _: integer, _: integer, _:integer)
 ---@field copy fun(self: TSTree): TSTree
 ---@field included_ranges fun(self: TSTree, include_bytes: true): Range6[]
 ---@field included_ranges fun(self: TSTree, include_bytes: false): Range4[]
 
+---@class TSQuery: userdata
+---@field inspect fun(self: TSQuery): TSQueryInfo
+
+---@class (exact) TSQueryInfo
+---@field captures string[]
+---@field patterns table<integer, (integer|string)[][]>
+
 ---@return integer
 vim._ts_get_language_version = function() end
 
 ---@return integer
 vim._ts_get_minimum_language_version = function() end
 
+---@param lang string Language to use for the query
+---@param query string Query string in s-expr syntax
+---@return TSQuery
+vim._ts_parse_query = function(lang, query) end
+
 ---@param lang string
 ---@return TSParser
 vim._create_ts_parser = function(lang) end
diff --git a/runtime/lua/vim/treesitter/_query_linter.lua b/runtime/lua/vim/treesitter/_query_linter.lua
index 87d74789a3..6216d4e891 100644
--- a/runtime/lua/vim/treesitter/_query_linter.lua
+++ b/runtime/lua/vim/treesitter/_query_linter.lua
@@ -17,7 +17,7 @@ local M = {}
 --- @field is_first_lang boolean Whether this is the first language of a linter run checking queries for multiple `langs`
 
 --- Adds a diagnostic for node in the query buffer
---- @param diagnostics Diagnostic[]
+--- @param diagnostics vim.Diagnostic[]
 --- @param range Range4
 --- @param lint string
 --- @param lang string?
@@ -45,7 +45,7 @@ local function guess_query_lang(buf)
 end
 
 --- @param buf integer
---- @param opts QueryLinterOpts|QueryLinterNormalizedOpts|nil
+--- @param opts vim.treesitter.query.lint.Opts|QueryLinterNormalizedOpts|nil
 --- @return QueryLinterNormalizedOpts
 local function normalize_opts(buf, opts)
   opts = opts or {}
@@ -92,7 +92,7 @@ local function get_error_entry(err, node)
     end_col = end_col + #underlined
   elseif msg:match('^Invalid') then
     -- Use the length of the problematic type/capture/field
-    end_col = end_col + #msg:match('"([^"]+)"')
+    end_col = end_col + #(msg:match('"([^"]+)"') or '')
   end
 
   return {
@@ -114,7 +114,7 @@ end
 --- @return vim.treesitter.ParseError?
 local parse = vim.func._memoize(hash_parse, function(node, buf, lang)
   local query_text = vim.treesitter.get_node_text(node, buf)
-  local ok, err = pcall(vim.treesitter.query.parse, lang, query_text) ---@type boolean|vim.treesitter.ParseError, string|Query
+  local ok, err = pcall(vim.treesitter.query.parse, lang, query_text) ---@type boolean|vim.treesitter.ParseError, string|vim.treesitter.Query
 
   if not ok and type(err) == 'string' then
     return get_error_entry(err, node)
@@ -122,28 +122,30 @@ local parse = vim.func._memoize(hash_parse, function(node, buf, lang)
 end)
 
 --- @param buf integer
---- @param match table<integer,TSNode>
---- @param query Query
+--- @param match vim.treesitter.query.TSMatch
+--- @param query vim.treesitter.Query
 --- @param lang_context QueryLinterLanguageContext
---- @param diagnostics Diagnostic[]
+--- @param diagnostics vim.Diagnostic[]
 local function lint_match(buf, match, query, lang_context, diagnostics)
   local lang = lang_context.lang
   local parser_info = lang_context.parser_info
 
-  for id, node in pairs(match) do
-    local cap_id = query.captures[id]
+  for id, nodes in pairs(match) do
+    for _, node in ipairs(nodes) do
+      local cap_id = query.captures[id]
 
-    -- perform language-independent checks only for first lang
-    if lang_context.is_first_lang and cap_id == 'error' then
-      local node_text = vim.treesitter.get_node_text(node, buf):gsub('\n', ' ')
-      add_lint_for_node(diagnostics, { node:range() }, 'Syntax error: ' .. node_text)
-    end
+      -- perform language-independent checks only for first lang
+      if lang_context.is_first_lang and cap_id == 'error' then
+        local node_text = vim.treesitter.get_node_text(node, buf):gsub('\n', ' ')
+        add_lint_for_node(diagnostics, { node:range() }, 'Syntax error: ' .. node_text)
+      end
 
-    -- other checks rely on Neovim parser introspection
-    if lang and parser_info and cap_id == 'toplevel' then
-      local err = parse(node, buf, lang)
-      if err then
-        add_lint_for_node(diagnostics, err.range, err.msg, lang)
+      -- other checks rely on Neovim parser introspection
+      if lang and parser_info and cap_id == 'toplevel' then
+        local err = parse(node, buf, lang)
+        if err then
+          add_lint_for_node(diagnostics, err.range, err.msg, lang)
+        end
       end
     end
   end
@@ -151,7 +153,7 @@ end
 
 --- @private
 --- @param buf integer Buffer to lint
---- @param opts QueryLinterOpts|QueryLinterNormalizedOpts|nil Options for linting
+--- @param opts vim.treesitter.query.lint.Opts|QueryLinterNormalizedOpts|nil Options for linting
 function M.lint(buf, opts)
   if buf == 0 then
     buf = api.nvim_get_current_buf()
@@ -173,7 +175,7 @@ function M.lint(buf, opts)
     parser:parse()
     parser:for_each_tree(function(tree, ltree)
       if ltree:lang() == 'query' then
-        for _, match, _ in query:iter_matches(tree:root(), buf, 0, -1) do
+        for _, match, _ in query:iter_matches(tree:root(), buf, 0, -1, { all = true }) do
           local lang_context = {
             lang = lang,
             parser_info = parser_info,
@@ -195,7 +197,7 @@ function M.clear(buf)
 end
 
 --- @private
---- @param findstart integer
+--- @param findstart 0|1
 --- @param base string
 function M.omnifunc(findstart, base)
   if findstart == 1 then
diff --git a/runtime/lua/vim/treesitter/dev.lua b/runtime/lua/vim/treesitter/dev.lua
index 69ddc9b558..dc2a14d238 100644
--- a/runtime/lua/vim/treesitter/dev.lua
+++ b/runtime/lua/vim/treesitter/dev.lua
@@ -1,31 +1,29 @@
 local api = vim.api
 
----@class TSDevModule
 local M = {}
 
----@class TSTreeView
+---@class (private) vim.treesitter.dev.TSTreeView
 ---@field ns integer API namespace
----@field opts table Options table with the following keys:
----                  - anon (boolean): If true, display anonymous nodes
----                  - lang (boolean): If true, display the language alongside each node
----                  - indent (number): Number of spaces to indent nested lines. Default is 2.
----@field nodes TSP.Node[]
----@field named TSP.Node[]
+---@field opts vim.treesitter.dev.TSTreeViewOpts
+---@field nodes vim.treesitter.dev.Node[]
+---@field named vim.treesitter.dev.Node[]
 local TSTreeView = {}
 
----@class TSP.Node
----@field id integer Node id
----@field text string Node text
----@field named boolean True if this is a named (non-anonymous) node
----@field depth integer Depth of the node within the tree
----@field lnum integer Beginning line number of this node in the source buffer
----@field col integer Beginning column number of this node in the source buffer
----@field end_lnum integer Final line number of this node in the source buffer
----@field end_col integer Final column number of this node in the source buffer
+---@private
+---@class (private) vim.treesitter.dev.TSTreeViewOpts
+---@field anon boolean If true, display anonymous nodes.
+---@field lang boolean If true, display the language alongside each node.
+---@field indent number Number of spaces to indent nested lines.
+
+---@class (private) vim.treesitter.dev.Node
+---@field node TSNode Treesitter node
+---@field field string? Node field
+---@field depth integer Depth of this node in the tree
+---@field text string? Text displayed in the inspector for this node. Not computed until the
+---                    inspector is drawn.
 ---@field lang string Source language of this node
----@field root TSNode
 
----@class TSP.Injection
+---@class (private) vim.treesitter.dev.Injection
 ---@field lang string Source language of this injection
 ---@field root TSNode Root node of the injection
 
@@ -43,48 +41,26 @@ local TSTreeView = {}
 ---
 ---@param node TSNode Starting node to begin traversal |tsnode|
 ---@param depth integer Current recursion depth
+---@param field string|nil The field of the current node
 ---@param lang string Language of the tree currently being traversed
----@param injections table<string, TSP.Injection> Mapping of node ids to root nodes
+---@param injections table<string, vim.treesitter.dev.Injection> Mapping of node ids to root nodes
 ---                  of injected language trees (see explanation above)
----@param tree TSP.Node[] Output table containing a list of tables each representing a node in the tree
-local function traverse(node, depth, lang, injections, tree)
+---@param tree vim.treesitter.dev.Node[] Output table containing a list of tables each representing a node in the tree
+local function traverse(node, depth, field, lang, injections, tree)
+  table.insert(tree, {
+    node = node,
+    depth = depth,
+    lang = lang,
+    field = field,
+  })
+
   local injection = injections[node:id()]
   if injection then
-    traverse(injection.root, depth, injection.lang, injections, tree)
+    traverse(injection.root, depth + 1, nil, injection.lang, injections, tree)
   end
 
-  for child, field in node:iter_children() do
-    local type = child:type()
-    local lnum, col, end_lnum, end_col = child:range()
-    local named = child:named()
-    local text ---@type string
-    if named then
-      if field then
-        text = string.format('%s: (%s', field, type)
-      else
-        text = string.format('(%s', type)
-      end
-    else
-      text = string.format('"%s"', type:gsub('\n', '\\n'):gsub('"', '\\"'))
-    end
-
-    table.insert(tree, {
-      id = child:id(),
-      text = text,
-      named = named,
-      depth = depth,
-      lnum = lnum,
-      col = col,
-      end_lnum = end_lnum,
-      end_col = end_col,
-      lang = lang,
-    })
-
-    traverse(child, depth + 1, lang, injections, tree)
-
-    if named then
-      tree[#tree].text = string.format('%s)', tree[#tree].text)
-    end
+  for child, child_field in node:iter_children() do
+    traverse(child, depth + 1, child_field, lang, injections, tree)
   end
 
   return tree
@@ -95,44 +71,45 @@ end
 ---@param bufnr integer Source buffer number
 ---@param lang string|nil Language of source buffer
 ---
----@return TSTreeView|nil
+---@return vim.treesitter.dev.TSTreeView|nil
 ---@return string|nil Error message, if any
 ---
 ---@package
 function TSTreeView:new(bufnr, lang)
   local ok, parser = pcall(vim.treesitter.get_parser, bufnr or 0, lang)
   if not ok then
-    return nil, 'No parser available for the given buffer'
+    local err = parser --[[ @as string ]]
+    return nil, 'No parser available for the given buffer:\n' .. err
   end
 
   -- For each child tree (injected language), find the root of the tree and locate the node within
   -- the primary tree that contains that root. Add a mapping from the node in the primary tree to
   -- the root in the child tree to the {injections} table.
   local root = parser:parse(true)[1]:root()
-  local injections = {} ---@type table<string, TSP.Injection>
+  local injections = {} ---@type table<string, vim.treesitter.dev.Injection>
 
   parser:for_each_tree(function(parent_tree, parent_ltree)
     local parent = parent_tree:root()
     for _, child in pairs(parent_ltree:children()) do
-      child:for_each_tree(function(tree, ltree)
+      for _, tree in pairs(child:trees()) do
         local r = tree:root()
         local node = assert(parent:named_descendant_for_range(r:range()))
         local id = node:id()
         if not injections[id] or r:byte_length() > injections[id].root:byte_length() then
           injections[id] = {
-            lang = ltree:lang(),
+            lang = child:lang(),
             root = r,
           }
         end
-      end)
+      end
     end
   end)
 
-  local nodes = traverse(root, 0, parser:lang(), injections, {})
+  local nodes = traverse(root, 0, nil, parser:lang(), injections, {})
 
-  local named = {} ---@type TSP.Node[]
+  local named = {} ---@type vim.treesitter.dev.Node[]
   for _, v in ipairs(nodes) do
-    if v.named then
+    if v.node:named() then
       named[#named + 1] = v
     end
   end
@@ -141,6 +118,7 @@ function TSTreeView:new(bufnr, lang)
     ns = api.nvim_create_namespace('treesitter/dev-inspect'),
     nodes = nodes,
     named = named,
+    ---@type vim.treesitter.dev.TSTreeViewOpts
     opts = {
       anon = false,
       lang = false,
@@ -155,16 +133,12 @@ end
 
 local decor_ns = api.nvim_create_namespace('ts.dev')
 
----@param lnum integer
----@param col integer
----@param end_lnum integer
----@param end_col integer
+---@param range Range4
 ---@return string
-local function get_range_str(lnum, col, end_lnum, end_col)
-  if lnum == end_lnum then
-    return string.format('[%d:%d - %d]', lnum + 1, col + 1, end_col)
-  end
-  return string.format('[%d:%d - %d:%d]', lnum + 1, col + 1, end_lnum + 1, end_col)
+local function range_to_string(range)
+  ---@type integer, integer, integer, integer
+  local row, col, end_row, end_col = unpack(range)
+  return string.format('[%d, %d] - [%d, %d]', row, col, end_row, end_col)
 end
 
 ---@param w integer
@@ -183,7 +157,10 @@ end
 local function set_dev_properties(w, b)
   vim.wo[w].scrolloff = 5
   vim.wo[w].wrap = false
-  vim.wo[w].foldmethod = 'manual' -- disable folding
+  vim.wo[w].foldmethod = 'expr'
+  vim.wo[w].foldexpr = 'v:lua.vim.treesitter.foldexpr()' -- explicitly set foldexpr
+  vim.wo[w].foldenable = false -- Don't fold on first open InspectTree
+  vim.wo[w].foldlevel = 99
   vim.bo[b].buflisted = false
   vim.bo[b].buftype = 'nofile'
   vim.bo[b].bufhidden = 'wipe'
@@ -192,7 +169,7 @@ end
 
 --- Updates the cursor position in the inspector to match the node under the cursor.
 ---
---- @param treeview TSTreeView
+--- @param treeview vim.treesitter.dev.TSTreeView
 --- @param lang string
 --- @param source_buf integer
 --- @param inspect_buf integer
@@ -213,7 +190,7 @@ local function set_inspector_cursor(treeview, lang, source_buf, inspect_buf, ins
 
   local cursor_node_id = cursor_node:id()
   for i, v in treeview:iter() do
-    if v.id == cursor_node_id then
+    if v.node:id() == cursor_node_id then
       local start = v.depth * treeview.opts.indent ---@type integer
       local end_col = start + #v.text
       api.nvim_buf_set_extmark(inspect_buf, treeview.ns, i - 1, start, {
@@ -228,6 +205,8 @@ end
 
 --- Write the contents of this View into {bufnr}.
 ---
+--- Calling this function computes the text that is displayed for each node.
+---
 ---@param bufnr integer Buffer number to write into.
 ---@package
 function TSTreeView:draw(bufnr)
@@ -235,13 +214,35 @@ function TSTreeView:draw(bufnr)
   local lines = {} ---@type string[]
   local lang_hl_marks = {} ---@type table[]
 
-  for _, item in self:iter() do
-    local range_str = get_range_str(item.lnum, item.col, item.end_lnum, item.end_col)
+  for i, item in self:iter() do
+    local range_str = range_to_string({ item.node:range() })
     local lang_str = self.opts.lang and string.format(' %s', item.lang) or ''
+
+    local text ---@type string
+    if item.node:named() then
+      if item.field then
+        text = string.format('%s: (%s', item.field, item.node:type())
+      else
+        text = string.format('(%s', item.node:type())
+      end
+    else
+      text = string.format('"%s"', item.node:type():gsu