diff options
Diffstat (limited to 'src/nvim/lua/treesitter.c')
-rw-r--r-- | src/nvim/lua/treesitter.c | 231 |
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; |