tree-sitter: inspect language

3 files changed, 53 insertions, 1 deletions

diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 1b5f416b67..3a1b1fc4b3 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -53,5 +53,9 @@ end
 
 -- TODO: weak table with reusable parser per buffer.
 
-return {create_parser=create_parser, add_language=vim._ts_add_language}
+return {
+  create_parser=create_parser,
+  add_language=vim._ts_add_language,
+  inspect_language=vim._ts_inspect_language,
+}
 
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index c208711985..aa83e3c1ba 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -841,4 +841,7 @@ static void nlua_add_treesitter(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
 
   lua_pushcfunction(lstate, ts_lua_register_lang);
   lua_setfield(lstate, -2, "_ts_add_language");
+
+  lua_pushcfunction(lstate, ts_lua_inspect_lang);
+  lua_setfield(lstate, -2, "_ts_inspect_language");
 }
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 9d599da85f..db337db533 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -157,6 +157,51 @@ int ts_lua_register_lang(lua_State *L)
   return 1;
 }
 
+int ts_lua_inspect_lang(lua_State *L)
+{
+  if (lua_gettop(L) < 1 || !lua_isstring(L, 1)) {
+    return luaL_error(L, "string expected");
+  }
+  const char *lang_name = lua_tostring(L, 1);
+
+  TSLanguage *lang = pmap_get(cstr_t)(langs, lang_name);
+  if (!lang) {
+    return luaL_error(L, "no such language: %s", lang_name);
+  }
+
+  lua_createtable(L, 0, 2);  // [retval]
+
+  size_t nsymbols = (size_t)ts_language_symbol_count(lang);
+
+  lua_createtable(L, nsymbols-1, 1);  // [retval, symbols]
+  for (size_t i = 0; i < nsymbols; i++) {
+    TSSymbolType t = ts_language_symbol_type(lang, i);
+    if (t == TSSymbolTypeAuxiliary) {
+      // not used by the API
+      continue;
+    }
+    lua_createtable(L, 2, 0);  // [retval, symbols, elem]
+    lua_pushstring(L, ts_language_symbol_name(lang, i));
+    lua_rawseti(L, -2, 1);
+    lua_pushboolean(L, t == TSSymbolTypeRegular);
+    lua_rawseti(L, -2, 2);  // [retval, symbols, elem]
+    lua_rawseti(L, -2, i);  // [retval, symbols]
+  }
+
+  lua_setfield(L, -2, "symbols");  // [retval]
+
+  // TODO: this seems to be empty, what langs have fields?
+  size_t nfields = (size_t)ts_language_field_count(lang);
+  lua_createtable(L, nfields-1, 1);  // [retval, fields]
+  for (size_t i = 0; i < nfields; i++) {
+    lua_pushstring(L, ts_language_field_name_for_id(lang, i));
+    lua_rawseti(L, -2, i);  // [retval, fields]
+  }
+
+  lua_setfield(L, -2, "fields");  // [retval]
+  return 1;
+}
+
 int tslua_push_parser(lua_State *L, const char *lang_name)
 {
   TSParser *parser = ts_parser_new();