1 files changed, 120 insertions, 3 deletions

diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index da64685a40..a9e7838980 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -20,6 +20,7 @@
 #include "nvim/api/private/helpers.h"
 #include "nvim/buffer_defs.h"
 #include "nvim/globals.h"
+#include "nvim/lua/executor.h"
 #include "nvim/lua/treesitter.h"
 #include "nvim/macros.h"
 #include "nvim/map.h"
@@ -43,6 +44,13 @@ typedef struct {
   int max_match_id;
 } TSLua_cursor;
 
+typedef struct {
+  LuaRef cb;
+  lua_State *lstate;
+  bool lex;
+  bool parse;
+} TSLuaLoggerOpts;
+
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "lua/treesitter.c.generated.h"
 #endif
@@ -56,6 +64,8 @@ static struct luaL_Reg parser_meta[] = {
   { "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 }
 };
 
@@ -231,9 +241,9 @@ int tslua_remove_lang(lua_State *L)
   const char *lang_name = luaL_checkstring(L, 1);
   bool present = pmap_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;
@@ -322,6 +332,12 @@ static int parser_gc(lua_State *L)
     return 0;
   }
 
+  TSLogger logger = ts_parser_logger(*p);
+  if (logger.log) {
+    TSLuaLoggerOpts *opts = (TSLuaLoggerOpts *)logger.payload;
+    xfree(opts);
+  }
+
   ts_parser_delete(*p);
   return 0;
 }
@@ -669,9 +685,82 @@ static int parser_get_timeout(lua_State *L)
   }
 
   lua_pushinteger(L, (long)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;
+
+  nlua_pushref(lstate, 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));
+
+  *opts = (TSLuaLoggerOpts){
+    .lex = lua_toboolean(L, 2),
+    .parse = lua_toboolean(L, 3),
+    .cb = nlua_ref_global(L, 4),
+    .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;
+    nlua_pushref(L, opts->cb);
+  } else {
+    lua_pushnil(L);
+  }
+
+  return 1;
+}
+
 // Tree methods
 
 /// push tree interface on lua stack.
@@ -1349,6 +1438,11 @@ static int node_rawquery(lua_State *L)
   } else {
     cursor = ts_query_cursor_new();
   }
+
+#ifdef NVIM_TS_HAS_SET_MAX_START_DEPTH
+  // reset the start depth
+  ts_query_cursor_set_max_start_depth(cursor, 0);
+#endif
   ts_query_cursor_set_match_limit(cursor, 256);
   ts_query_cursor_exec(cursor, query, node);
 
@@ -1360,6 +1454,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;