1 files changed, 416 insertions, 101 deletions

diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 56f4daed1a..008b3f2e95 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -1,11 +1,9 @@
-// This is an open source non-commercial project. Dear PVS-Studio, please check
-// it. PVS-Studio Static Code Analyzer for C, C++ and C#: http://www.viva64.com
-
 // lua bindings for tree-sitter.
 // NB: this file mostly contains a generic lua interface for tree-sitter
 // trees and nodes, and could be broken out as a reusable lua package
 
 #include <assert.h>
+#include <ctype.h>
 #include <lauxlib.h>
 #include <limits.h>
 #include <lua.h>
@@ -14,6 +12,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <tree_sitter/api.h>
 #include <uv.h>
 
 #include "klib/kvec.h"
@@ -21,14 +20,13 @@
 #include "nvim/buffer_defs.h"
 #include "nvim/globals.h"
 #include "nvim/lua/treesitter.h"
-#include "nvim/macros.h"
-#include "nvim/map.h"
+#include "nvim/macros_defs.h"
+#include "nvim/map_defs.h"
 #include "nvim/memline.h"
 #include "nvim/memory.h"
-#include "nvim/pos.h"
+#include "nvim/pos_defs.h"
 #include "nvim/strings.h"
-#include "nvim/types.h"
-#include "tree_sitter/api.h"
+#include "nvim/types_defs.h"
 
 #define TS_META_PARSER "treesitter_parser"
 #define TS_META_TREE "treesitter_tree"
@@ -43,6 +41,17 @@ typedef struct {
   int max_match_id;
 } TSLua_cursor;
 
+typedef struct {
+  LuaRef cb;
+  lua_State *lstate;
+  bool lex;
+  bool parse;
+} TSLuaLoggerOpts;
+
+typedef struct {
+  TSTree *tree;
+} TSLuaTree;
+
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "lua/treesitter.c.generated.h"
 #endif
@@ -51,8 +60,13 @@ static struct luaL_Reg parser_meta[] = {
   { "__gc", parser_gc },
   { "__tostring", parser_tostring },
   { "parse", parser_parse },
+  { "reset", parser_reset },
   { "set_included_ranges", parser_set_ranges },
   { "included_ranges", parser_get_ranges },
+  { "set_timeout", parser_set_timeout },
+  { "timeout", parser_get_timeout },
+  { "_set_logger", parser_set_logger },
+  { "_logger", parser_get_logger },
   { NULL, NULL }
 };
 
@@ -61,6 +75,7 @@ static struct luaL_Reg tree_meta[] = {
   { "__tostring", tree_tostring },
   { "root", tree_root },
   { "edit", tree_edit },
+  { "included_ranges", tree_get_ranges },
   { "copy", tree_copy },
   { NULL, NULL }
 };
@@ -78,6 +93,8 @@ static struct luaL_Reg node_meta[] = {
   { "field", node_field },
   { "named", node_named },
   { "missing", node_missing },
+  { "extra", node_extra },
+  { "has_changes", node_has_changes },
   { "has_error", node_has_error },
   { "sexpr", node_sexpr },
   { "child_count", node_child_count },
@@ -95,7 +112,9 @@ static struct luaL_Reg node_meta[] = {
   { "prev_named_sibling", node_prev_named_sibling },
   { "named_children", node_named_children },
   { "root", node_root },
+  { "tree", node_tree },
   { "byte_length", node_byte_length },
+  { "equal", node_equal },
 
   { NULL, NULL }
 };
@@ -132,9 +151,9 @@ static void build_meta(lua_State *L, const char *tname, const luaL_Reg *meta)
   lua_pop(L, 1);  // [] (don't use it now)
 }
 
-/// init the tslua library
+/// Init the tslua library.
 ///
-/// all global state is stored in the regirstry of the lua_State
+/// All global state is stored in the registry of the lua_State.
 void tslua_init(lua_State *L)
 {
   // type metatables
@@ -145,15 +164,13 @@ void tslua_init(lua_State *L)
   build_meta(L, TS_META_QUERYCURSOR, querycursor_meta);
   build_meta(L, TS_META_TREECURSOR, treecursor_meta);
 
-#ifdef NVIM_TS_HAS_SET_ALLOCATOR
   ts_set_allocator(xmalloc, xcalloc, xrealloc, xfree);
-#endif
 }
 
 int tslua_has_language(lua_State *L)
 {
   const char *lang_name = luaL_checkstring(L, 1);
-  lua_pushboolean(L, pmap_has(cstr_t)(&langs, lang_name));
+  lua_pushboolean(L, map_has(cstr_t, &langs, lang_name));
   return 1;
 }
 
@@ -170,7 +187,7 @@ int tslua_add_language(lua_State *L)
     symbol_name = luaL_checkstring(L, 3);
   }
 
-  if (pmap_has(cstr_t)(&langs, lang_name)) {
+  if (map_has(cstr_t, &langs, lang_name)) {
     lua_pushboolean(L, true);
     return 1;
   }
@@ -223,11 +240,11 @@ int tslua_add_language(lua_State *L)
 int tslua_remove_lang(lua_State *L)
 {
   const char *lang_name = luaL_checkstring(L, 1);
-  bool present = pmap_has(cstr_t)(&langs, lang_name);
+  bool present = map_has(cstr_t, &langs, lang_name);
   if (present) {
-    char *key = (char *)pmap_key(cstr_t)(&langs, lang_name);
-    pmap_del(cstr_t)(&langs, lang_name);
-    xfree(key);
+    cstr_t key;
+    pmap_del(cstr_t)(&langs, lang_name, &key);
+    xfree((void *)key);
   }
   lua_pushboolean(L, present);
   return 1;
@@ -309,6 +326,17 @@ static TSParser **parser_check(lua_State *L, uint16_t index)
   return luaL_checkudata(L, index, TS_META_PARSER);
 }
 
+static void logger_gc(TSLogger logger)
+{
+  if (!logger.log) {
+    return;
+  }
+
+  TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)logger.payload;
+  luaL_unref(opts->lstate, LUA_REGISTRYINDEX, opts->cb);
+  xfree(opts);
+}
+
 static int parser_gc(lua_State *L)
 {
   TSParser **p = parser_check(L, 1);
@@ -316,6 +344,7 @@ static int parser_gc(lua_State *L)
     return 0;
   }
 
+  logger_gc(ts_parser_logger(*p));
   ts_parser_delete(*p);
   return 0;
 }
@@ -329,7 +358,7 @@ static int parser_tostring(lua_State *L)
 static const char *input_cb(void *payload, uint32_t byte_index, TSPoint position,
                             uint32_t *bytes_read)
 {
-  buf_T *bp  = payload;
+  buf_T *bp = payload;
 #define BUFSIZE 256
   static char buf[BUFSIZE];
 
@@ -337,7 +366,7 @@ static const char *input_cb(void *payload, uint32_t byte_index, TSPoint position
     *bytes_read = 0;
     return "";
   }
-  char *line = ml_get_buf(bp, (linenr_T)position.row + 1, false);
+  char *line = ml_get_buf(bp, (linenr_T)position.row + 1);
   size_t len = strlen(line);
   if (position.column > len) {
     *bytes_read = 0;
@@ -359,19 +388,29 @@ static const char *input_cb(void *payload, uint32_t byte_index, TSPoint position
 #undef BUFSIZE
 }
 
-static void push_ranges(lua_State *L, const TSRange *ranges, const size_t length)
+static void push_ranges(lua_State *L, const TSRange *ranges, const size_t length,
+                        bool include_bytes)
 {
   lua_createtable(L, (int)length, 0);
   for (size_t i = 0; i < length; i++) {
-    lua_createtable(L, 4, 0);
+    lua_createtable(L, include_bytes ? 6 : 4, 0);
+    int j = 1;
     lua_pushinteger(L, ranges[i].start_point.row);
-    lua_rawseti(L, -2, 1);
+    lua_rawseti(L, -2, j++);
     lua_pushinteger(L, ranges[i].start_point.column);
-    lua_rawseti(L, -2, 2);
+    lua_rawseti(L, -2, j++);
+    if (include_bytes) {
+      lua_pushinteger(L, ranges[i].start_byte);
+      lua_rawseti(L, -2, j++);
+    }
     lua_pushinteger(L, ranges[i].end_point.row);
-    lua_rawseti(L, -2, 3);
+    lua_rawseti(L, -2, j++);
     lua_pushinteger(L, ranges[i].end_point.column);
-    lua_rawseti(L, -2, 4);
+    lua_rawseti(L, -2, j++);
+    if (include_bytes) {
+      lua_pushinteger(L, ranges[i].end_byte);
+      lua_rawseti(L, -2, j++);
+    }
 
     lua_rawseti(L, -2, (int)(i + 1));
   }
@@ -386,14 +425,14 @@ static int parser_parse(lua_State *L)
 
   TSTree *old_tree = NULL;
   if (!lua_isnil(L, 2)) {
-    TSTree **tmp = tree_check(L, 2);
-    old_tree = tmp ? *tmp : NULL;
+    TSLuaTree *ud = tree_check(L, 2);
+    old_tree = ud ? ud->tree : NULL;
   }
 
   TSTree *new_tree = NULL;
   size_t len;
   const char *str;
-  long bufnr;
+  handle_T bufnr;
   buf_T *buf;
   TSInput input;
 
@@ -406,13 +445,13 @@ static int parser_parse(lua_State *L)
     break;
 
   case LUA_TNUMBER:
-    bufnr = lua_tointeger(L, 3);
-    buf = handle_get_buffer((handle_T)bufnr);
+    bufnr = (handle_T)lua_tointeger(L, 3);
+    buf = handle_get_buffer(bufnr);
 
     if (!buf) {
 #define BUFSIZE 256
       char ebuf[BUFSIZE] = { 0 };
-      vim_snprintf(ebuf, BUFSIZE, "invalid buffer handle: %ld", bufnr);
+      vim_snprintf(ebuf, BUFSIZE, "invalid buffer handle: %d", bufnr);
       return luaL_argerror(L, 3, ebuf);
 #undef BUFSIZE
     }
@@ -426,34 +465,46 @@ static int parser_parse(lua_State *L)
     return luaL_argerror(L, 3, "expected either string or buffer handle");
   }
 
+  bool include_bytes = (lua_gettop(L) >= 4) && lua_toboolean(L, 4);
+
   // Sometimes parsing fails (timeout, or wrong parser ABI)
   // In those case, just return an error.
   if (!new_tree) {
     return luaL_error(L, "An error occurred when parsing.");
   }
 
-  // The new tree will be pushed to the stack, without copy, ownership is now to
-  // the lua GC.
-  // Old tree is still owned by the lua GC.
+  // The new tree will be pushed to the stack, without copy, ownership is now to the lua GC.
+  // Old tree is owned by lua GC since before
   uint32_t n_ranges = 0;
-  TSRange *changed = old_tree ?  ts_tree_get_changed_ranges(old_tree, new_tree, &n_ranges) : NULL;
+  TSRange *changed = old_tree ? ts_tree_get_changed_ranges(old_tree, new_tree, &n_ranges) : NULL;
 
-  push_tree(L, new_tree, false);  // [tree]
+  push_tree(L, new_tree);  // [tree]
 
-  push_ranges(L, changed, n_ranges);  // [tree, ranges]
+  push_ranges(L, changed, n_ranges, include_bytes);  // [tree, ranges]
 
   xfree(changed);
   return 2;
 }
 
+static int parser_reset(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (p && *p) {
+    ts_parser_reset(*p);
+  }
+
+  return 0;
+}
+
 static int tree_copy(lua_State *L)
 {
-  TSTree **tree = tree_check(L, 1);
-  if (!(*tree)) {
+  TSLuaTree *ud = tree_check(L, 1);
+  if (!ud) {
     return 0;
   }
 
-  push_tree(L, *tree, true);  // [tree]
+  TSTree *copy = ts_tree_copy(ud->tree);
+  push_tree(L, copy);  // [tree]
 
   return 1;
 }
@@ -465,8 +516,8 @@ static int tree_edit(lua_State *L)
     return lua_error(L);
   }
 
-  TSTree **tree = tree_check(L, 1);
-  if (!(*tree)) {
+  TSLuaTree *ud = tree_check(L, 1);
+  if (!ud) {
     return 0;
   }
 
@@ -480,11 +531,29 @@ static int tree_edit(lua_State *L)
   TSInputEdit edit = { start_byte, old_end_byte, new_end_byte,
                        start_point, old_end_point, new_end_point };
 
-  ts_tree_edit(*tree, &edit);
+  ts_tree_edit(ud->tree, &edit);
 
   return 0;
 }
 
+static int tree_get_ranges(lua_State *L)
+{
+  TSLuaTree *ud = tree_check(L, 1);
+  if (!ud) {
+    return 0;
+  }
+
+  bool include_bytes = (lua_gettop(L) >= 2) && lua_toboolean(L, 2);
+
+  uint32_t len;
+  TSRange *ranges = ts_tree_included_ranges(ud->tree, &len);
+
+  push_ranges(L, ranges, len, include_bytes);
+
+  xfree(ranges);
+  return 1;
+}
+
 // Use the top of the stack (without popping it) to create a TSRange, it can be
 // either a lua table or a TSNode
 static void range_from_lua(lua_State *L, TSRange *range)
@@ -590,58 +659,159 @@ static int parser_get_ranges(lua_State *L)
     return 0;
   }
 
+  bool include_bytes = (lua_gettop(L) >= 2) && lua_toboolean(L, 2);
+
   uint32_t len;
   const TSRange *ranges = ts_parser_included_ranges(*p, &len);
 
-  push_ranges(L, ranges, len);
+  push_ranges(L, ranges, len, include_bytes);
+  return 1;
+}
+
+static int parser_set_timeout(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (!p) {
+    return 0;
+  }
+
+  if (lua_gettop(L) < 2) {
+    luaL_error(L, "integer expected");
+  }
+
+  uint32_t timeout = (uint32_t)luaL_checkinteger(L, 2);
+  ts_parser_set_timeout_micros(*p, timeout);
+  return 0;
+}
+
+static int parser_get_timeout(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (!p) {
+    return 0;
+  }
+
+  lua_pushinteger(L, (lua_Integer)ts_parser_timeout_micros(*p));
+  return 1;
+}
+
+static void logger_cb(void *payload, TSLogType logtype, const char *s)
+{
+  TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)payload;
+  if ((!opts->lex && logtype == TSLogTypeLex)
+      || (!opts->parse && logtype == TSLogTypeParse)) {
+    return;
+  }
+
+  lua_State *lstate = opts->lstate;
+
+  lua_rawgeti(lstate, LUA_REGISTRYINDEX, opts->cb);
+  lua_pushstring(lstate, logtype == TSLogTypeParse ? "parse" : "lex");
+  lua_pushstring(lstate, s);
+  if (lua_pcall(lstate, 2, 0, 0)) {
+    luaL_error(lstate, "Error executing treesitter logger callback");
+  }
+}
+
+static int parser_set_logger(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (!p) {
+    return 0;
+  }
+
+  if (!lua_isboolean(L, 2)) {
+    return luaL_argerror(L, 2, "boolean expected");
+  }
+
+  if (!lua_isboolean(L, 3)) {
+    return luaL_argerror(L, 3, "boolean expected");
+  }
+
+  if (!lua_isfunction(L, 4)) {
+    return luaL_argerror(L, 4, "function expected");
+  }
+
+  TSLuaLoggerOpts *opts = xmalloc(sizeof(TSLuaLoggerOpts));
+  lua_pushvalue(L, 4);
+  LuaRef ref = luaL_ref(L, LUA_REGISTRYINDEX);
+
+  *opts = (TSLuaLoggerOpts){
+    .lex = lua_toboolean(L, 2),
+    .parse = lua_toboolean(L, 3),
+    .cb = ref,
+    .lstate = L
+  };
+
+  TSLogger logger = {
+    .payload = (void *)opts,
+    .log = logger_cb
+  };
+
+  ts_parser_set_logger(*p, logger);
+  return 0;
+}
+
+static int parser_get_logger(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (!p) {
+    return 0;
+  }
+
+  TSLogger logger = ts_parser_logger(*p);
+  if (logger.log) {
+    TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)logger.payload;
+    lua_rawgeti(L, LUA_REGISTRYINDEX, opts->cb);
+  } else {
+    lua_pushnil(L);
+  }
+
   return 1;
 }
 
 // Tree methods
 
-/// push tree interface on lua stack.
+/// Push tree interface on to the lua stack.
 ///
-/// This makes a copy of the tree, so ownership of the argument is unaffected.
-void push_tree(lua_State *L, TSTree *tree, bool do_copy)
+/// The tree is not copied. Ownership of the tree is transferred from C to
+/// Lua. If needed use ts_tree_copy() in the caller
+static void push_tree(lua_State *L, TSTree *tree)
 {
   if (tree == NULL) {
     lua_pushnil(L);
     return;
   }
-  TSTree **ud = lua_newuserdata(L, sizeof(TSTree *));  // [udata]
 
-  if (do_copy) {
-    *ud = ts_tree_copy(tree);
-  } else {
-    *ud = tree;
-  }
+  TSLuaTree *ud = lua_newuserdata(L, sizeof(TSLuaTree));  // [udata]
+
+  ud->tree = tree;
 
   lua_getfield(L, LUA_REGISTRYINDEX, TS_META_TREE);  // [udata, meta]
   lua_setmetatable(L, -2);  // [udata]
 
-  // table used for node wrappers to keep a reference to tree wrapper
-  // NB: in lua 5.3 the uservalue for the node could just be the tree, but
-  // in lua 5.1 the uservalue (fenv) must be a table.
+  // To prevent the tree from being garbage collected, create a reference to it
+  // in the fenv which will be passed to userdata nodes of the tree.
+  // Note: environments (fenvs) associated with userdata have no meaning in Lua
+  // and are only used to associate a table.
   lua_createtable(L, 1, 0);  // [udata, reftable]
   lua_pushvalue(L, -2);  // [udata, reftable, udata]
   lua_rawseti(L, -2, 1);  // [udata, reftable]
   lua_setfenv(L, -2);  // [udata]
 }
 
-static TSTree **tree_check(lua_State *L, int index)
+static TSLuaTree *tree_check(lua_State *L, int index)
 {
-  TSTree **ud = luaL_checkudata(L, index, TS_META_TREE);
+  TSLuaTree *ud = luaL_checkudata(L, index, TS_META_TREE);
   return ud;
 }
 
 static int tree_gc(lua_State *L)
 {
-  TSTree **tree = tree_check(L, 1);
-  if (!tree) {
-    return 0;
+  TSLuaTree *ud = tree_check(L, 1);
+  if (ud) {
+    ts_tree_delete(ud->tree);
   }
-
-  ts_tree_delete(*tree);
   return 0;
 }
 
@@ -653,20 +823,20 @@ static int tree_tostring(lua_State *L)
 
 static int tree_root(lua_State *L)
 {
-  TSTree **tree = tree_check(L, 1);
-  if (!tree) {
+  TSLuaTree *ud = tree_check(L, 1);
+  if (!ud) {
     return 0;
   }
-  TSNode root = ts_tree_root_node(*tree);
+  TSNode root = ts_tree_root_node(ud->tree);
   push_node(L, root, 1);
   return 1;
 }
 
 // Node methods
 
-/// push node interface on lua stack
+/// Push node interface on to the Lua stack
 ///
-/// top of stack must either be the tree this node belongs to or another node
+/// Top of stack must either be the tree this node belongs to or another node
 /// of the same tree! This value is not popped. Can only be called inside a
 /// cfunction with the tslua environment.
 static void push_node(lua_State *L, TSNode node, int uindex)
@@ -680,6 +850,8 @@ static void push_node(lua_State *L, TSNode node, int uindex)
   *ud = node;
   lua_getfield(L, LUA_REGISTRYINDEX, TS_META_NODE);  // [udata, meta]
   lua_setmetatable(L, -2);  // [udata]
+
+  // Copy the fenv which contains the nodes tree.
   lua_getfenv(L, uindex);  // [udata, reftable]
   lua_setfenv(L, -2);  // [udata]
 }
@@ -740,12 +912,26 @@ static int node_range(lua_State *L)
   if (!node_check(L, 1, &node)) {
     return 0;
   }
+
+  bool include_bytes = (lua_gettop(L) >= 2) && lua_toboolean(L, 2);
+
   TSPoint start = ts_node_start_point(node);
   TSPoint end = ts_node_end_point(node);
-  lua_pushnumber(L, start.row);
-  lua_pushnumber(L, start.column);
-  lua_pushnumber(L, end.row);
-  lua_pushnumber(L, end.column);
+
+  if (include_bytes) {
+    lua_pushinteger(L, start.row);
+    lua_pushinteger(L, start.column);
+    lua_pushinteger(L, ts_node_start_byte(node));
+    lua_pushinteger(L, end.row);
+    lua_pushinteger(L, end.column);
+    lua_pushinteger(L, ts_node_end_byte(node));
+    return 6;
+  }
+
+  lua_pushinteger(L, start.row);
+  lua_pushinteger(L, start.column);
+  lua_pushinteger(L, end.row);
+  lua_pushinteger(L, end.column);
   return 4;
 }
 
@@ -757,9 +943,9 @@ static int node_start(lua_State *L)
   }
   TSPoint start = ts_node_start_point(node);
   uint32_t start_byte = ts_node_start_byte(node);
-  lua_pushnumber(L, start.row);
-  lua_pushnumber(L, start.column);
-  lua_pushnumber(L, start_byte);
+  lua_pushinteger(L, start.row);
+  lua_pushinteger(L, start.column);
+  lua_pushinteger(L, start_byte);
   return 3;
 }
 
@@ -771,9 +957,9 @@ static int node_end(lua_State *L)
   }
   TSPoint end = ts_node_end_point(node);
   uint32_t end_byte = ts_node_end_byte(node);
-  lua_pushnumber(L, end.row);
-  lua_pushnumber(L, end.column);
-  lua_pushnumber(L, end_byte);
+  lua_pushinteger(L, end.row);
+  lua_pushinteger(L, end.column);
+  lua_pushinteger(L, end_byte);
   return 3;
 }
 
@@ -784,7 +970,7 @@ static int node_child_count(lua_State *L)
     return 0;
   }
   uint32_t count = ts_node_child_count(node);
-  lua_pushnumber(L, count);
+  lua_pushinteger(L, count);
   return 1;
 }
 
@@ -795,7 +981,7 @@ static int node_named_child_count(lua_State *L)
     return 0;
   }
   uint32_t count = ts_node_named_child_count(node);
-  lua_pushnumber(L, count);
+  lua_pushinteger(L, count);
   return 1;
 }
 
@@ -816,7 +1002,7 @@ static int node_symbol(lua_State *L)
     return 0;
   }
   TSSymbol symbol = ts_node_symbol(node);
-  lua_pushnumber(L, symbol);
+  lua_pushinteger(L, symbol);
   return 1;
 }
 
@@ -882,6 +1068,26 @@ static int node_missing(lua_State *L)
   return 1;
 }
 
+static int node_extra(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, 1, &node)) {
+    return 0;
+  }
+  lua_pushboolean(L, ts_node_is_extra(node));
+  return 1;
+}
+
+static int node_has_changes(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, 1, &node)) {
+    return 0;
+  }
+  lua_pushboolean(L, ts_node_has_changes(node));
+  return 1;
+}
+
 static int node_has_error(lua_State *L)
 {
   TSNode node;
@@ -898,8 +1104,8 @@ static int node_child(lua_State *L)
   if (!node_check(L, 1, &node)) {
     return 0;
   }
-  long num = lua_tointeger(L, 2);
-  TSNode child = ts_node_child(node, (uint32_t)num);
+  uint32_t num = (uint32_t)lua_tointeger(L, 2);
+  TSNode child = ts_node_child(node, num);
 
   push_node(L, child, 1);
   return 1;
@@ -911,8 +1117,8 @@ static int node_named_child(lua_State *L)
   if (!node_check(L, 1, &node)) {
     return 0;
   }
-  long num = lua_tointeger(L, 2);
-  TSNode child = ts_node_named_child(node, (uint32_t)num);
+  uint32_t num = (uint32_t)lua_tointeger(L, 2);
+  TSNode child = ts_node_named_child(node, num);
 
   push_node(L, child, 1);
   return 1;
@@ -1108,6 +1314,19 @@ static int node_root(lua_State *L)
   return 1;
 }
 
+static int node_tree(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, 1, &node)) {
+    return 0;
+  }
+
+  lua_getfenv(L, 1);  // [udata, reftable]
+  lua_rawgeti(L, -1, 1);  // [udata, reftable, tree_udata]
+
+  return 1;
+}
+
 static int node_byte_length(lua_State *L)
 {
   TSNode node;
@@ -1118,7 +1337,23 @@ static int node_byte_length(lua_State *L)
   uint32_t start_byte = ts_node_start_byte(node);
   uint32_t end_byte = ts_node_end_byte(node);
 
-  lua_pushnumber(L, end_byte - start_byte);
+  lua_pushinteger(L, end_byte - start_byte);
+  return 1;
+}
+
+static int node_equal(lua_State *L)
+{
+  TSNode node1;
+  if (!node_check(L, 1, &node1)) {
+    return 0;
+  }
+
+  TSNode node2;
+  if (!node_check(L, 2, &node2)) {
+    return luaL_error(L, "TSNode expected");
+  }
+
+  lua_pushboolean(L, ts_node_eq(node1, node2));
   return 1;
 }
 
@@ -1213,11 +1448,12 @@ static int node_rawquery(lua_State *L)
   } else {
     cursor = ts_query_cursor_new();
   }
-  // TODO(clason): API introduced after tree-sitter release 0.19.5
-  // remove guard when minimum ts version is bumped to 0.19.6+
-#ifdef NVIM_TS_HAS_SET_MATCH_LIMIT
-  ts_query_cursor_set_match_limit(cursor, 64);
+
+#ifdef NVIM_TS_HAS_SET_MAX_START_DEPTH
+  // reset the start depth
+  ts_query_cursor_set_max_start_depth(cursor, UINT32_MAX);
 #endif
+  ts_query_cursor_set_match_limit(cursor, 256);
   ts_query_cursor_exec(cursor, query, node);
 
   bool captures = lua_toboolean(L, 3);
@@ -1228,6 +1464,29 @@ static int node_rawquery(lua_State *L)
     ts_query_cursor_set_point_range(cursor, (TSPoint){ start, 0 }, (TSPoint){ end, 0 });
   }
 
+  if (lua_gettop(L) >= 6 && !lua_isnil(L, 6)) {
+    if (!lua_istable(L, 6)) {
+      return luaL_error(L, "table expected");
+    }
+    lua_pushnil(L);
+    // stack: [dict, ..., nil]
+    while (lua_next(L, 6)) {
+      // stack: [dict, ..., key, value]
+      if (lua_type(L, -2) == LUA_TSTRING) {
+        char *k = (char *)lua_tostring(L, -2);
+        if (strequal("max_start_depth", k)) {
+          // TODO(lewis6991): remove ifdef when min TS version is 0.20.9
+#ifdef NVIM_TS_HAS_SET_MAX_START_DEPTH
+          uint32_t max_start_depth = (uint32_t)lua_tointeger(L, -1);
+          ts_query_cursor_set_max_start_depth(cursor, max_start_depth);
+#endif
+        }
+      }
+      lua_pop(L, 1);  // pop the value; lua_next will pop the key.
+      // stack: [dict, ..., key]
+    }
+  }
+
   TSLua_cursor *ud = lua_newuserdata(L, sizeof(*ud));  // [udata]
   ud->cursor = cursor;
   ud->predicated_match = -1;
@@ -1281,8 +1540,9 @@ int tslua_parse_query(lua_State *L)
   TSQuery *query = ts_query_new(lang, src, (uint32_t)len, &error_offset, &error_type);
 
   if (!query) {
-    return luaL_error(L, "query: %s at position %d for language %s",
-                      query_err_string(error_type), (int)error_offset, lang_name);
+    char err_msg[IOSIZE];
+    query_err_string(src, (int)error_offset, error_type, err_msg, sizeof(err_msg));
+    return luaL_error(L, "%s", err_msg);
   }
 
   TSQuery **ud = lua_newuserdata(L, sizeof(TSQuery *));  // [udata]
@@ -1292,24 +1552,79 @@ int tslua_parse_query(lua_State *L)
   return 1;
 }
 
-static const char *query_err_string(TSQueryError err)
+static const char *query_err_to_string(TSQueryError error_type)
 {
-  switch (err) {
+  switch (error_type) {
   case TSQueryErrorSyntax:
-    return "invalid syntax";
+    return "Invalid syntax:\n";
   case TSQueryErrorNodeType:
-    return "invalid node type";
+    return "Invalid node type ";
   case TSQueryErrorField:
-    return "invalid field";
+    return "Invalid field name ";
   case TSQueryErrorCapture:
-    return "invalid capture";
+    return "Invalid capture name ";
   case TSQueryErrorStructure:
-    return "invalid structure";
+    return "Impossible pattern:\n";
   default:
     return "error";
   }
 }
 
+static void query_err_string(const char *src, int error_offset, TSQueryError error_type, char *err,
+                             size_t errlen)
+{
+  int line_start = 0;
+  int row = 0;
+  const char *error_line = NULL;
+  int error_line_len = 0;
+
+  const char *end_str;
+  do {
+    const char *src_tmp = src + line_start;
+    end_str = strchr(src_tmp, '\n');
+    int line_length = end_str != NULL ? (int)(end_str - src_tmp) : (int)strlen(src_tmp);
+    int line_end = line_start + line_length;
+    if (line_end > error_offset) {
+      error_line = src_tmp;
+      error_line_len = line_length;
+      break;
+    }
+    line_start = line_end + 1;
+    row++;
+  } while (end_str != NULL);
+
+  int column = error_offset - line_start;
+
+  const char *type_msg = query_err_to_string(error_type);
+  snprintf(err, errlen, "Query error at %d:%d. %s", row + 1, column + 1, type_msg);
+  size_t offset = strlen(err);
+  errlen = errlen - offset;
+  err = err + offset;
+
+  // Error types that report names
+  if (error_type == TSQueryErrorNodeType
+      || error_type == TSQueryErrorField
+      || error_type == TSQueryErrorCapture) {
+    const char *suffix = src + error_offset;
+    int suffix_len = 0;
+    char c = suffix[suffix_len];
+    while (isalnum(c) || c == '_' || c == '-' || c == '.') {
+      c = suffix[++suffix_len];
+    }
+    snprintf(err, errlen, "\"%.*s\":\n", suffix_len, suffix);
+    offset = strlen(err);
+    errlen = errlen - offset;
+    err = err + offset;
+  }
+
+  if (!error_line) {
+    snprintf(err, errlen, "Unexpected EOF\n");
+    return;
+  }
+
+  snprintf(err, errlen, "%.*s\n%*s^\n", error_line_len, error_line, column, "");
+}
+
 static TSQuery *query_check(lua_State *L, int index)
 {
   TSQuery **ud = luaL_checkudata(L, index, TS_META_QUERY);
@@ -1367,7 +1682,7 @@ static int query_inspect(lua_State *L)
                                                        &strlen);
         lua_pushlstring(L, str, strlen);  // [retval, patterns, pat, pred, item]
       } else if (step[k].type == TSQueryPredicateStepTypeCapture) {
-        lua_pushnumber(L, step[k].value_id + 1);  // [..., pat, pred, item]
+        lua_pushinteger(L, step[k].value_id + 1);  // [..., pat, pred, item]
       } else {
         abort();
       }