1 files changed, 194 insertions, 37 deletions

diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 138031237e..5258352e72 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -39,10 +39,11 @@ typedef struct {
 static struct luaL_Reg parser_meta[] = {
   { "__gc", parser_gc },
   { "__tostring", parser_tostring },
-  { "parse_buf", parser_parse_buf },
+  { "parse", parser_parse },
   { "edit", parser_edit },
   { "tree", parser_tree },
   { "set_included_ranges", parser_set_ranges },
+  { "included_ranges", parser_get_ranges },
   { NULL, NULL }
 };
 
@@ -62,6 +63,7 @@ static struct luaL_Reg node_meta[] = {
   { "end_", node_end },
   { "type", node_type },
   { "symbol", node_symbol },
+  { "field", node_field },
   { "named", node_named },
   { "missing", node_missing },
   { "has_error", node_has_error },
@@ -73,6 +75,7 @@ static struct luaL_Reg node_meta[] = {
   { "descendant_for_range", node_descendant_for_range },
   { "named_descendant_for_range", node_named_descendant_for_range },
   { "parent", node_parent },
+  { "iter_children", node_iter_children },
   { "_rawquery", node_rawquery },
   { NULL, NULL }
 };
@@ -84,12 +87,17 @@ static struct luaL_Reg query_meta[] = {
   { NULL, NULL }
 };
 
-// cursor is not exposed, but still needs garbage collection
+// cursors are not exposed, but still needs garbage collection
 static struct luaL_Reg querycursor_meta[] = {
   { "__gc", querycursor_gc },
   { NULL, NULL }
 };
 
+static struct luaL_Reg treecursor_meta[] = {
+  { "__gc", treecursor_gc },
+  { NULL, NULL }
+};
+
 static PMap(cstr_t) *langs;
 
 static void build_meta(lua_State *L, const char *tname, const luaL_Reg *meta)
@@ -116,6 +124,7 @@ void tslua_init(lua_State *L)
   build_meta(L, "treesitter_node", node_meta);
   build_meta(L, "treesitter_query", query_meta);
   build_meta(L, "treesitter_querycursor", querycursor_meta);
+  build_meta(L, "treesitter_treecursor", treecursor_meta);
 }
 
 int tslua_has_language(lua_State *L)
@@ -166,6 +175,14 @@ int tslua_add_language(lua_State *L)
     return luaL_error(L, "Failed to load parser: internal error");
   }
 
+  uint32_t lang_version = ts_language_version(lang);
+  if (lang_version < TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION) {
+    return luaL_error(
+        L,
+        "ABI version mismatch : expected %" PRIu32 ", found %" PRIu32,
+        TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION, lang_version);
+  }
+
   pmap_put(cstr_t)(langs, xstrdup(lang_name), lang);
 
   lua_pushboolean(L, true);
@@ -278,43 +295,84 @@ static const char *input_cb(void *payload, uint32_t byte_index,
   }
   char_u *line = ml_get_buf(bp, position.row+1, false);
   size_t len = STRLEN(line);
-
   if (position.column > len) {
     *bytes_read = 0;
-  } else {
-    size_t tocopy = MIN(len-position.column, BUFSIZE);
-
-    memcpy(buf, line+position.column, tocopy);
-    // Translate embedded \n to NUL
-    memchrsub(buf, '\n', '\0', tocopy);
-    *bytes_read = (uint32_t)tocopy;
-    if (tocopy < BUFSIZE) {
-      // now add the final \n. If it didn't fit, input_cb will be called again
-      // on the same line with advanced column.
-      buf[tocopy] = '\n';
-      (*bytes_read)++;
-    }
+    return "";
+  }
+  size_t tocopy = MIN(len-position.column, BUFSIZE);
+
+  memcpy(buf, line+position.column, tocopy);
+  // Translate embedded \n to NUL
+  memchrsub(buf, '\n', '\0', tocopy);
+  *bytes_read = (uint32_t)tocopy;
+  if (tocopy < BUFSIZE) {
+    // now add the final \n. If it didn't fit, input_cb will be called again
+    // on the same line with advanced column.
+    buf[tocopy] = '\n';
+    (*bytes_read)++;
   }
   return buf;
 #undef BUFSIZE
 }
 
-static int parser_parse_buf(lua_State *L)
+static void push_ranges(lua_State *L,
+                        const TSRange *ranges,
+                        const unsigned int length)
+{
+  lua_createtable(L, length, 0);
+  for (size_t i = 0; i < length; i++) {
+    lua_createtable(L, 4, 0);
+    lua_pushinteger(L, ranges[i].start_point.row);
+    lua_rawseti(L, -2, 1);
+    lua_pushinteger(L, ranges[i].start_point.column);
+    lua_rawseti(L, -2, 2);
+    lua_pushinteger(L, ranges[i].end_point.row);
+    lua_rawseti(L, -2, 3);
+    lua_pushinteger(L, ranges[i].end_point.column);
+    lua_rawseti(L, -2, 4);
+
+    lua_rawseti(L, -2, i+1);
+  }
+}
+
+static int parser_parse(lua_State *L)
 {
   TSLua_parser *p = parser_check(L);
   if (!p) {
     return 0;
   }
 
-  long bufnr = lua_tointeger(L, 2);
-  buf_T *buf = handle_get_buffer(bufnr);
+  TSTree *new_tree;
+  size_t len;
+  const char *str;
+  long bufnr;
+  buf_T *buf;
+  TSInput input;
+
+  // This switch is necessary because of the behavior of lua_isstring, that
+  // consider numbers as strings...
+  switch (lua_type(L, 2)) {
+    case LUA_TSTRING:
+      str = lua_tolstring(L, 2, &len);
+      new_tree = ts_parser_parse_string(p->parser, p->tree, str, len);
+      break;
+
+    case LUA_TNUMBER:
+      bufnr = lua_tointeger(L, 2);
+      buf = handle_get_buffer(bufnr);
+
+      if (!buf) {
+        return luaL_error(L, "invalid buffer handle: %d", bufnr);
+      }
 
-  if (!buf) {
-    return luaL_error(L, "invalid buffer handle: %d", bufnr);
-  }
+      input = (TSInput){ (void *)buf, input_cb, TSInputEncodingUTF8 };
+      new_tree = ts_parser_parse(p->parser, p->tree, input);
 
-  TSInput input = { (void *)buf, input_cb, TSInputEncodingUTF8 };
-  TSTree *new_tree = ts_parser_parse(p->parser, p->tree, input);
+      break;
+
+    default:
+      return luaL_error(L, "invalid argument to parser:parse()");
+  }
 
   uint32_t n_ranges = 0;
   TSRange *changed = p->tree ? ts_tree_get_changed_ranges(p->tree, new_tree,
@@ -326,20 +384,8 @@ static int parser_parse_buf(lua_State *L)
 
   tslua_push_tree(L, p->tree);
 
-  lua_createtable(L, n_ranges, 0);
-  for (size_t i = 0; i < n_ranges; i++) {
-    lua_createtable(L, 4, 0);
-    lua_pushinteger(L, changed[i].start_point.row);
-    lua_rawseti(L, -2, 1);
-    lua_pushinteger(L, changed[i].start_point.column);
-    lua_rawseti(L, -2, 2);
-    lua_pushinteger(L, changed[i].end_point.row);
-    lua_rawseti(L, -2, 3);
-    lua_pushinteger(L, changed[i].end_point.column);
-    lua_rawseti(L, -2, 4);
+  push_ranges(L, changed, n_ranges);
 
-    lua_rawseti(L, -2, i+1);
-  }
   xfree(changed);
   return 2;
 }
@@ -437,6 +483,21 @@ static int parser_set_ranges(lua_State *L)
   return 0;
 }
 
+static int parser_get_ranges(lua_State *L)
+{
+  TSLua_parser *p = parser_check(L);
+  if (!p || !p->parser) {
+    return 0;
+  }
+
+  unsigned int len;
+  const TSRange *ranges = ts_parser_included_ranges(p->parser, &len);
+
+  push_ranges(L, ranges, len);
+
+  return 1;
+}
+
 
 // Tree methods
 
@@ -646,6 +707,34 @@ static int node_symbol(lua_State *L)
   return 1;
 }
 
+static int node_field(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, 1, &node)) {
+    return 0;
+  }
+
+  size_t name_len;
+  const char *field_name = luaL_checklstring(L, 2, &name_len);
+
+  TSTreeCursor cursor = ts_tree_cursor_new(node);
+
+  lua_newtable(L);  // [table]
+  unsigned int curr_index = 0;
+
+  if (ts_tree_cursor_goto_first_child(&cursor)) {
+    do {
+      if (!STRCMP(field_name, ts_tree_cursor_current_field_name(&cursor))) {
+        push_node(L, ts_tree_cursor_current_node(&cursor), 1);  // [table, node]
+        lua_rawseti(L, -2, ++curr_index);
+      }
+    } while (ts_tree_cursor_goto_next_sibling(&cursor));
+  }
+
+  ts_tree_cursor_delete(&cursor);
+  return 1;
+}
+
 static int node_named(lua_State *L)
 {
   TSNode node;
@@ -746,6 +835,74 @@ static int node_named_descendant_for_range(lua_State *L)
   return 1;
 }
 
+static int node_next_child(lua_State *L)
+{
+  TSTreeCursor *ud = luaL_checkudata(
+      L, lua_upvalueindex(1), "treesitter_treecursor");
+  if (!ud) {
+    return 0;
+  }
+
+  TSNode source;
+  if (!node_check(L, lua_upvalueindex(2), &source)) {
+    return 0;
+  }
+
+  // First call should return first child
+  if (ts_node_eq(source, ts_tree_cursor_current_node(ud))) {
+    if (ts_tree_cursor_goto_first_child(ud)) {
+      goto push;
+    } else {
+      goto end;
+    }
+  }
+
+  if (ts_tree_cursor_goto_next_sibling(ud)) {
+push:
+      push_node(
+          L,
+          ts_tree_cursor_current_node(ud),
+          lua_upvalueindex(2));  // [node]
+
+      const char * field = ts_tree_cursor_current_field_name(ud);
+
+      if (field != NULL) {
+        lua_pushstring(L, ts_tree_cursor_current_field_name(ud));
+      } else {
+        lua_pushnil(L);
+      }  // [node, field_name_or_nil]
+      return 2;
+  }
+
+end:
+  return 0;
+}
+
+static int node_iter_children(lua_State *L)
+{
+  TSNode source;
+  if (!node_check(L, 1, &source)) {
+    return 0;
+  }
+
+  TSTreeCursor *ud = lua_newuserdata(L, sizeof(TSTreeCursor));  // [udata]
+  *ud = ts_tree_cursor_new(source);
+
+  lua_getfield(L, LUA_REGISTRYINDEX, "treesitter_treecursor");  // [udata, mt]
+  lua_setmetatable(L, -2);  // [udata]
+  lua_pushvalue(L, 1);  // [udata, source_node]
+  lua_pushcclosure(L, node_next_child, 2);
+
+  return 1;
+}
+
+static int treecursor_gc(lua_State *L)
+{
+  TSTreeCursor *ud = luaL_checkudata(L, 1, "treesitter_treecursor");
+  ts_tree_cursor_delete(ud);
+  return 0;
+}
+
 static int node_parent(lua_State *L)
 {
   TSNode node;