1 files changed, 48 insertions, 29 deletions

diff --git a/runtime/lua/vim/treesitter/query.lua b/runtime/lua/vim/treesitter/query.lua
index 66da179ea3..8383551b5f 100644
--- a/runtime/lua/vim/treesitter/query.lua
+++ b/runtime/lua/vim/treesitter/query.lua
@@ -48,7 +48,7 @@ function M.get_query_files(lang, query_name, is_included)
   local base_langs = {}
 
   -- Now get the base languages by looking at the first line of every file
-  -- The syntax is the folowing :
+  -- The syntax is the following :
   -- ;+ inherits: ({language},)*{language}
   --
   -- {language} ::= {lang} | ({lang})
@@ -138,48 +138,77 @@ function M.get_query(lang, query_name)
   end
 end
 
+local query_cache = setmetatable({}, {
+  __index = function(tbl, key)
+    rawset(tbl, key, {})
+    return rawget(tbl, key)
+  end
+})
+
 --- Parse {query} as a string. (If the query is in a file, the caller
----       should read the contents into a string before calling).
+--- should read the contents into a string before calling).
 ---
 --- Returns a `Query` (see |lua-treesitter-query|) object which can be used to
 --- search nodes in the syntax tree for the patterns defined in {query}
 --- using `iter_*` methods below.
 ---
---- Exposes `info` and `captures` with additional information about the {query}.
+--- Exposes `info` and `captures` with additional context about {query}.
 ---   - `captures` contains the list of unique capture names defined in
 ---     {query}.
 ---   -` info.captures` also points to `captures`.
 ---   - `info.patterns` contains information about predicates.
 ---
----@param lang The language
----@param query A string containing the query (s-expr syntax)
+---@param lang string The language
+---@param query string A string containing the query (s-expr syntax)
 ---
 ---@returns The query
 function M.parse_query(lang, query)
   language.require_language(lang)
-  local self = setmetatable({}, Query)
-  self.query = vim._ts_parse_query(lang, query)
-  self.info = self.query:inspect()
-  self.captures = self.info.captures
-  return self
+  local cached = query_cache[lang][query]
+  if cached then
+    return cached
+  else
+    local self = setmetatable({}, Query)
+    self.query = vim._ts_parse_query(lang, query)
+    self.info = self.query:inspect()
+    self.captures = self.info.captures
+    query_cache[lang][query] = self
+    return self
+  end
 end
 
--- TODO(vigoux): support multiline nodes too
-
 --- Gets the text corresponding to a given node
 ---
 ---@param node the node
----@param bsource The buffer or string from which the node is extracted
+---@param source The buffer or string from which the node is extracted
 function M.get_node_text(node, source)
   local start_row, start_col, start_byte = node:start()
   local end_row, end_col, end_byte = node:end_()
 
   if type(source) == "number" then
-    if start_row ~= end_row then
+    local lines
+    local eof_row = a.nvim_buf_line_count(source)
+    if start_row >= eof_row then
       return nil
     end
-    local line = a.nvim_buf_get_lines(source, start_row, start_row+1, true)[1]
-    return string.sub(line, start_col+1, end_col)
+
+    if end_col == 0 then
+      lines = a.nvim_buf_get_lines(source, start_row, end_row, true)
+      end_col = -1
+    else
+      lines = a.nvim_buf_get_lines(source, start_row, end_row + 1, true)
+    end
+
+    if #lines > 0 then
+      if #lines == 1 then
+        lines[1]      = string.sub(lines[1], start_col+1, end_col)
+      else
+        lines[1]      = string.sub(lines[1], start_col+1)
+        lines[#lines] = string.sub(lines[#lines], 1, end_col)
+      end
+    end
+
+    return table.concat(lines, "\n")
   elseif type(source) == "string" then
     return source:sub(start_byte+1, end_byte)
   end
@@ -211,11 +240,6 @@ local predicate_handlers = {
   ["lua-match?"] = function(match, _, source, predicate)
       local node = match[predicate[2]]
       local regex = predicate[3]
-      local start_row, _, end_row, _ = node:range()
-      if start_row ~= end_row then
-        return false
-      end
-
       return string.find(M.get_node_text(node, source), regex)
   end,
 
@@ -239,13 +263,8 @@ local predicate_handlers = {
 
     return function(match, _, source, pred)
       local node = match[pred[2]]
-      local start_row, start_col, end_row, end_col = node:range()
-      if start_row ~= end_row then
-        return false
-      end
-
       local regex = compiled_vim_regexes[pred[3]]
-      return regex:match_line(source, start_row, start_col, end_col)
+      return regex:match_str(M.get_node_text(node, source))
     end
   end)(),
 
@@ -446,7 +465,7 @@ end
 ---
 --- {source} is needed if the query contains predicates, then the caller
 --- must ensure to use a freshly parsed tree consistent with the current
---- text of the buffer (if relevent). {start_row} and {end_row} can be used to limit
+--- text of the buffer (if relevant). {start_row} and {end_row} can be used to limit
 --- matches inside a row range (this is typically used with root node
 --- as the node, i e to get syntax highlight matches in the current
 --- viewport). When omitted the start and end row values are used from the given node.
@@ -466,7 +485,7 @@ end
 --- </pre>
 ---
 ---@param node The node under which the search will occur
----@param source The source buffer or string to exctract text from
+---@param source The source buffer or string to extract text from
 ---@param start The starting line of the search
 ---@param stop The stopping line of the search (end-exclusive)
 ---