tree-sitter: implement query functionality and highlighting prototype [skip.lint]

13 files changed, 578 insertions, 47 deletions

diff --git a/src/nvim/api/buffer.c b/src/nvim/api/buffer.c
index 8f5718d97e..e6f8f73b9d 100644
--- a/src/nvim/api/buffer.c
+++ b/src/nvim/api/buffer.c
@@ -169,21 +169,21 @@ Boolean nvim_buf_attach(uint64_t channel_id,
         goto error;
       }
       cb.on_lines = v->data.luaref;
-      v->data.integer = LUA_NOREF;
+      v->data.luaref = LUA_NOREF;
     } else if (is_lua && strequal("on_changedtick", k.data)) {
       if (v->type != kObjectTypeLuaRef) {
         api_set_error(err, kErrorTypeValidation, "callback is not a function");
         goto error;
       }
       cb.on_changedtick = v->data.luaref;
-      v->data.integer = LUA_NOREF;
+      v->data.luaref = LUA_NOREF;
     } else if (is_lua && strequal("on_detach", k.data)) {
       if (v->type != kObjectTypeLuaRef) {
         api_set_error(err, kErrorTypeValidation, "callback is not a function");
         goto error;
       }
       cb.on_detach = v->data.luaref;
-      v->data.integer = LUA_NOREF;
+      v->data.luaref = LUA_NOREF;
     } else if (is_lua && strequal("utf_sizes", k.data)) {
       if (v->type != kObjectTypeBoolean) {
         api_set_error(err, kErrorTypeValidation, "utf_sizes must be boolean");
@@ -231,6 +231,90 @@ Boolean nvim_buf_detach(uint64_t channel_id,
   return true;
 }
 
+static void buf_clear_luahl(buf_T *buf, bool force)
+{
+  if (buf->b_luahl || force) {
+    executor_free_luaref(buf->b_luahl_start);
+    executor_free_luaref(buf->b_luahl_window);
+    executor_free_luaref(buf->b_luahl_line);
+    executor_free_luaref(buf->b_luahl_end);
+  }
+  buf->b_luahl_start = LUA_NOREF;
+  buf->b_luahl_window = LUA_NOREF;
+  buf->b_luahl_line = LUA_NOREF;
+  buf->b_luahl_end = LUA_NOREF;
+}
+
+/// Unstabilized interface for defining syntax hl in lua.
+///
+/// This is not yet safe for general use, lua callbacks will need to
+/// be restricted, like textlock and probably other stuff.
+///
+/// The API on_line/nvim__put_attr is quite raw and not intended to be the
+/// final shape. Ideally this should operate on chunks larger than a single
+/// line to reduce interpreter overhead, and generate annotation objects
+/// (bufhl/virttext) on the fly but using the same representation.
+void nvim__buf_set_luahl(uint64_t channel_id, Buffer buffer,
+                         DictionaryOf(LuaRef) opts, Error *err)
+  FUNC_API_LUA_ONLY
+{
+  buf_T *buf = find_buffer_by_handle(buffer, err);
+
+  if (!buf) {
+    return;
+  }
+
+  redraw_buf_later(buf, NOT_VALID);
+  buf_clear_luahl(buf, false);
+
+  for (size_t i = 0; i < opts.size; i++) {
+    String k = opts.items[i].key;
+    Object *v = &opts.items[i].value;
+    if (strequal("on_start", k.data)) {
+      if (v->type != kObjectTypeLuaRef) {
+        api_set_error(err, kErrorTypeValidation, "callback is not a function");
+        goto error;
+      }
+      buf->b_luahl_start = v->data.luaref;
+      v->data.luaref = LUA_NOREF;
+    } else if (strequal("on_window", k.data)) {
+      if (v->type != kObjectTypeLuaRef) {
+        api_set_error(err, kErrorTypeValidation, "callback is not a function");
+        goto error;
+      }
+      buf->b_luahl_window = v->data.luaref;
+      v->data.luaref = LUA_NOREF;
+    } else if (strequal("on_line", k.data)) {
+      if (v->type != kObjectTypeLuaRef) {
+        api_set_error(err, kErrorTypeValidation, "callback is not a function");
+        goto error;
+      }
+      buf->b_luahl_line = v->data.luaref;
+      v->data.luaref = LUA_NOREF;
+    } else {
+      api_set_error(err, kErrorTypeValidation, "unexpected key: %s", k.data);
+      goto error;
+    }
+  }
+  buf->b_luahl = true;
+  return;
+error:
+  buf_clear_luahl(buf, true);
+  buf->b_luahl = false;
+}
+
+void nvim__buf_redraw_range(Buffer buffer, Integer first, Integer last,
+                            Error *err)
+  FUNC_API_LUA_ONLY
+{
+  buf_T *buf = find_buffer_by_handle(buffer, err);
+  if (!buf) {
+    return;
+  }
+
+  redraw_buf_range_later(buf, (linenr_T)first+1, (linenr_T)last);
+}
+
 /// Sets a buffer line
 ///
 /// @deprecated use nvim_buf_set_lines instead.
@@ -1112,7 +1196,6 @@ Array nvim_buf_get_extmarks(Buffer buffer, Integer ns_id, Object start,
         return rv;
       }
       limit = v->data.integer;
-      v->data.integer = LUA_NOREF;
     } else {
       api_set_error(err, kErrorTypeValidation, "unexpected key: %s", k.data);
       return rv;
diff --git a/src/nvim/api/private/helpers.h b/src/nvim/api/private/helpers.h
index 8930f252f6..048b937136 100644
--- a/src/nvim/api/private/helpers.h
+++ b/src/nvim/api/private/helpers.h
@@ -60,6 +60,12 @@
 #define ADD(array, item) \
   kv_push(array, item)
 
+#define FIXED_TEMP_ARRAY(name, fixsize) \
+  Array name = ARRAY_DICT_INIT; \
+  Object name##__items[fixsize]; \
+  args.size = fixsize; \
+  args.items = name##__items; \
+
 #define STATIC_CSTR_AS_STRING(s) ((String) {.data = s, .size = sizeof(s) - 1})
 
 /// Create a new String instance, putting data in allocated memory
diff --git a/src/nvim/api/vim.c b/src/nvim/api/vim.c
index 19601b6539..2f59edee33 100644
--- a/src/nvim/api/vim.c
+++ b/src/nvim/api/vim.c
@@ -189,6 +189,15 @@ Dictionary nvim_get_hl_by_id(Integer hl_id, Boolean rgb, Error *err)
   return hl_get_attr_by_id(attrcode, rgb, err);
 }
 
+/// Gets a highlight group by name
+///
+/// similar to |hlID()|, but allocates a new ID if not present.
+Integer nvim_get_hl_id_by_name(String name)
+  FUNC_API_SINCE(7)
+{
+  return syn_check_group((const char_u *)name.data, (int)name.size);
+}
+
 /// Sends input-keys to Nvim, subject to various quirks controlled by `mode`
 /// flags. This is a blocking call, unlike |nvim_input()|.
 ///
@@ -2546,3 +2555,27 @@ Array nvim__inspect_cell(Integer grid, Integer row, Integer col, Error *err)
   }
   return ret;
 }
+
+/// Set attrs in nvim__buf_set_lua_hl callbacks
+///
+/// TODO(bfredl): This is rather pedestrian. The final
+/// interface should probably be derived from a reformed
+/// bufhl/virttext interface with full support for multi-line
+/// ranges etc
+void nvim__put_attr(Integer id, Integer c0, Integer c1)
+  FUNC_API_LUA_ONLY
+{
+  if (!lua_attr_active) {
+    return;
+  }
+  if (id == 0 || syn_get_final_id((int)id) == 0) {
+    return;
+  }
+  int attr = syn_id2attr((int)id);
+  c0 = MAX(c0, 0);
+  c1 = MIN(c1, (Integer)lua_attr_bufsize);
+  for (Integer c = c0; c < c1; c++) {
+    lua_attr_buf[c] = (sattr_T)hl_combine_attr(lua_attr_buf[c], (int)attr);
+  }
+  return;
+}
diff --git a/src/nvim/buffer_defs.h b/src/nvim/buffer_defs.h
index 700d8b82e6..bc4fb2997d 100644
--- a/src/nvim/buffer_defs.h
+++ b/src/nvim/buffer_defs.h
@@ -832,6 +832,12 @@ struct file_buffer {
   // The number for times the current line has been flushed in the memline.
   int flush_count;
 
+  bool b_luahl;
+  LuaRef b_luahl_start;
+  LuaRef b_luahl_window;
+  LuaRef b_luahl_line;
+  LuaRef b_luahl_end;
+
   int b_diff_failed;    // internal diff failed for this buffer
 };
 
diff --git a/src/nvim/buffer_updates.c b/src/nvim/buffer_updates.c
index 3604578b50..d12527d6ac 100644
--- a/src/nvim/buffer_updates.c
+++ b/src/nvim/buffer_updates.c
@@ -157,7 +157,7 @@ void buf_updates_unregister_all(buf_T *buf)
       args.items[0] = BUFFER_OBJ(buf->handle);
 
       textlock++;
-      executor_exec_lua_cb(cb.on_detach, "detach", args, false);
+      executor_exec_lua_cb(cb.on_detach, "detach", args, false, NULL);
       textlock--;
     }
     free_update_callbacks(cb);
@@ -265,7 +265,7 @@ void buf_updates_send_changes(buf_T *buf,
         args.items[7] = INTEGER_OBJ((Integer)deleted_codeunits);
       }
       textlock++;
-      Object res = executor_exec_lua_cb(cb.on_lines, "lines", args, true);
+      Object res = executor_exec_lua_cb(cb.on_lines, "lines", args, true, NULL);
       textlock--;
 
       if (res.type == kObjectTypeBoolean && res.data.boolean == true) {
@@ -293,10 +293,7 @@ void buf_updates_changedtick(buf_T *buf)
     BufUpdateCallbacks cb = kv_A(buf->update_callbacks, i);
     bool keep = true;
     if (cb.on_changedtick != LUA_NOREF) {
-      Array args = ARRAY_DICT_INIT;
-      Object items[2];
-      args.size = 2;
-      args.items = items;
+      FIXED_TEMP_ARRAY(args, 2);
 
       // the first argument is always the buffer handle
       args.items[0] = BUFFER_OBJ(buf->handle);
@@ -306,7 +303,7 @@ void buf_updates_changedtick(buf_T *buf)
 
       textlock++;
       Object res = executor_exec_lua_cb(cb.on_changedtick, "changedtick",
-                                        args, true);
+                                        args, true, NULL);
       textlock--;
 
       if (res.type == kObjectTypeBoolean && res.data.boolean == true) {
diff --git a/src/nvim/func_attr.h b/src/nvim/func_attr.h
index 14b8158c7f..38b51b5564 100644
--- a/src/nvim/func_attr.h
+++ b/src/nvim/func_attr.h
@@ -211,6 +211,8 @@
 # define FUNC_API_NOEXPORT
 /// API function not exposed in VimL/eval.
 # define FUNC_API_REMOTE_ONLY
+/// API function not exposed in VimL/remote.
+# define FUNC_API_LUA_ONLY
 /// API function introduced at the given API level.
 # define FUNC_API_SINCE(X)
 /// API function deprecated since the given API level.
diff --git a/src/nvim/generators/c_grammar.lua b/src/nvim/generators/c_grammar.lua
index 1aa8223da0..de098b7a7b 100644
--- a/src/nvim/generators/c_grammar.lua
+++ b/src/nvim/generators/c_grammar.lua
@@ -42,6 +42,7 @@ local c_proto = Ct(
   (fill * Cg((P('FUNC_API_FAST') * Cc(true)), 'fast') ^ -1) *
   (fill * Cg((P('FUNC_API_NOEXPORT') * Cc(true)), 'noexport') ^ -1) *
   (fill * Cg((P('FUNC_API_REMOTE_ONLY') * Cc(true)), 'remote_only') ^ -1) *
+  (fill * Cg((P('FUNC_API_LUA_ONLY') * Cc(true)), 'lua_only') ^ -1) *
   (fill * Cg((P('FUNC_API_REMOTE_IMPL') * Cc(true)), 'remote_impl') ^ -1) *
   (fill * Cg((P('FUNC_API_BRIDGE_IMPL') * Cc(true)), 'bridge_impl') ^ -1) *
   (fill * Cg((P('FUNC_API_COMPOSITOR_IMPL') * Cc(true)), 'compositor_impl') ^ -1) *
diff --git a/src/nvim/generators/gen_api_dispatch.lua b/src/nvim/generators/gen_api_dispatch.lua
index 76dcf849d1..e861cfda35 100644
--- a/src/nvim/generators/gen_api_dispatch.lua
+++ b/src/nvim/generators/gen_api_dispatch.lua
@@ -192,7 +192,7 @@ end
 -- the real API.
 for i = 1, #functions do
   local fn = functions[i]
-  if fn.impl_name == nil then
+  if fn.impl_name == nil and not fn.lua_only then
     local args = {}
 
     output:write('Object handle_'..fn.name..'(uint64_t channel_id, Array args, Error *error)')
@@ -310,12 +310,13 @@ void msgpack_rpc_init_method_table(void)
 
 for i = 1, #functions do
   local fn = functions[i]
-  output:write('  msgpack_rpc_add_method_handler('..
-               '(String) {.data = "'..fn.name..'", '..
-               '.size = sizeof("'..fn.name..'") - 1}, '..
-               '(MsgpackRpcRequestHandler) {.fn = handle_'..  (fn.impl_name or fn.name)..
-               ', .fast = '..tostring(fn.fast)..'});\n')
-
+  if not fn.lua_only then
+      output:write('  msgpack_rpc_add_method_handler('..
+                   '(String) {.data = "'..fn.name..'", '..
+                   '.size = sizeof("'..fn.name..'") - 1}, '..
+                   '(MsgpackRpcRequestHandler) {.fn = handle_'..  (fn.impl_name or fn.name)..
+                   ', .fast = '..tostring(fn.fast)..'});\n')
+  end
 end
 
 output:write('\n}\n\n')
diff --git a/src/nvim/generators/gen_eval.lua b/src/nvim/generators/gen_eval.lua
index 2c6f8f2603..d16453530f 100644
--- a/src/nvim/generators/gen_eval.lua
+++ b/src/nvim/generators/gen_eval.lua
@@ -25,7 +25,7 @@ local gperfpipe = io.open(funcsfname .. '.gperf', 'wb')
 local funcs = require('eval').funcs
 local metadata = mpack.unpack(io.open(metadata_file, 'rb'):read("*all"))
 for _,fun in ipairs(metadata) do
-  if not fun.remote_only then
+  if not (fun.remote_only or fun.lua_only) then
     funcs[fun.name] = {
       args=#fun.parameters,
       func='api_wrapper',
diff --git a/src/nvim/globals.h b/src/nvim/globals.h
index 172c190df2..0a7a2d551e 100644
--- a/src/nvim/globals.h
+++ b/src/nvim/globals.h
@@ -126,6 +126,13 @@ typedef off_t off_T;
  */
 EXTERN int mod_mask INIT(= 0x0);                /* current key modifiers */
 
+
+// TODO(bfredl): for the final interface this should find a more suitable
+// location.
+EXTERN sattr_T *lua_attr_buf INIT(= NULL);
+EXTERN size_t lua_attr_bufsize INIT(= 0);
+EXTERN bool lua_attr_active INIT(= false);
+
 /*
  * Cmdline_row is the row where the command line starts, just below the
  * last window.
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 25f4be1c4d..1d3d9929d3 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -835,7 +835,7 @@ Object executor_exec_lua_api(const String str, const Array args, Error *err)
 }
 
 Object executor_exec_lua_cb(LuaRef ref, const char *name, Array args,
-                            bool retval)
+                            bool retval, Error *err)
 {
   lua_State *const lstate = nlua_enter();
   nlua_pushref(lstate, ref);
@@ -845,16 +845,24 @@ Object executor_exec_lua_cb(LuaRef ref, const char *name, Array args,
   }
 
   if (lua_pcall(lstate, (int)args.size+1, retval ? 1 : 0, 0)) {
-    // TODO(bfredl): callbacks:s might not always be msg-safe, for instance
-    // lua callbacks for redraw events. Later on let the caller deal with the
-    // error instead.
-    nlua_error(lstate, _("Error executing lua callback: %.*s"));
+    // if err is passed, the caller will deal with the error.
+    if (err) {
+      size_t len;
+      const char *errstr = lua_tolstring(lstate, -1, &len);
+      api_set_error(err, kErrorTypeException,
+                    "Error executing lua: %.*s", (int)len, errstr);
+    } else {
+      nlua_error(lstate, _("Error executing lua callback: %.*s"));
+    }
     return NIL;
   }
-  Error err = ERROR_INIT;
 
   if (retval) {
-    return nlua_pop_Object(lstate, false, &err);
+    Error dummy = ERROR_INIT;
+    if (err == NULL) {
+      err = &dummy;
+    }
+    return nlua_pop_Object(lstate, false, err);
   } else {
     return NIL;
   }
@@ -1007,4 +1015,7 @@ static void nlua_add_treesitter(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
 
   lua_pushcfunction(lstate, tslua_inspect_lang);
   lua_setfield(lstate, -2, "_ts_inspect_language");
+
+  lua_pushcfunction(lstate, ts_lua_parse_query);
+  lua_setfield(lstate, -2, "_ts_parse_query");
 }
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index d2072402bb..874fabd89f 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -26,6 +26,11 @@ typedef struct {
   TSTree *tree;  // internal tree, used for editing/reparsing
 } TSLua_parser;
 
+typedef struct {
+  TSQueryCursor *cursor;
+  int predicated_match;
+} TSLua_cursor;
+
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "lua/treesitter.c.generated.h"
 #endif
@@ -66,6 +71,20 @@ 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 },
+  { "_rawquery", node_rawquery },
+  { NULL, NULL }
+};
+
+static struct luaL_Reg query_meta[] = {
+  { "__gc", query_gc },
+  { "__tostring", query_tostring },
+  { "inspect", query_inspect },
+  { NULL, NULL }
+};
+
+// cursor is not exposed, but still needs garbage collection
+static struct luaL_Reg querycursor_meta[] = {
+  { "__gc", querycursor_gc },
   { NULL, NULL }
 };
 
@@ -96,6 +115,8 @@ void tslua_init(lua_State *L)
   build_meta(L, "treesitter_parser", parser_meta);
   build_meta(L, "treesitter_tree", tree_meta);
   build_meta(L, "treesitter_node", node_meta);
+  build_meta(L, "treesitter_query", query_meta);
+  build_meta(L, "treesitter_querycursor", querycursor_meta);
 }
 
 int tslua_register_lang(lua_State *L)
@@ -276,13 +297,33 @@ static int parser_parse_buf(lua_State *L)
   }
   TSInput input = { payload, input_cb, TSInputEncodingUTF8 };
   TSTree *new_tree = ts_parser_parse(p->parser, p->tree, input);
+
+  uint32_t n_ranges = 0;
+  TSRange *changed = p->tree ? ts_tree_get_changed_ranges(p->tree, new_tree,
+                                                          &n_ranges) : NULL;
   if (p->tree) {
     ts_tree_delete(p->tree);
   }
   p->tree = new_tree;
 
   tslua_push_tree(L, p->tree);
-  return 1;
+
+  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);
+
+    lua_rawseti(L, -2, i+1);
+  }
+  xfree(changed);
+  return 2;
 }
 
 static int parser_tree(lua_State *L)
@@ -383,7 +424,7 @@ static int tree_root(lua_State *L)
     return 0;
   }
   TSNode root = ts_tree_root_node(tree);
-  push_node(L, root);
+  push_node(L, root, 1);
   return 1;
 }
 
@@ -394,18 +435,19 @@ static int tree_root(lua_State *L)
 /// 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)
+static void push_node(lua_State *L, TSNode node, int uindex)
 {
+  assert(uindex > 0 || uindex < -LUA_MINSTACK);
   if (ts_node_is_null(node)) {
-    lua_pushnil(L);  // [src, nil]
+    lua_pushnil(L);  // [nil]
     return;
   }
-  TSNode *ud = lua_newuserdata(L, sizeof(TSNode));  // [src, udata]
+  TSNode *ud = lua_newuserdata(L, sizeof(TSNode));  // [udata]
   *ud = node;
-  lua_getfield(L, LUA_REGISTRYINDEX, "treesitter_node");  // [src, udata, meta]
-  lua_setmetatable(L, -2);  // [src, udata]
-  lua_getfenv(L, -2);  // [src, udata, reftable]
-  lua_setfenv(L, -2);  // [src, udata]
+  lua_getfield(L, LUA_REGISTRYINDEX, "treesitter_node");  // [udata, meta]
+  lua_setmetatable(L, -2);  // [udata]
+  lua_getfenv(L, uindex);  // [udata, reftable]
+  lua_setfenv(L, -2);  // [udata]
 }
 
 static bool node_check(lua_State *L, TSNode *res)
@@ -586,8 +628,7 @@ static int node_child(lua_State *L)
   long num = lua_tointeger(L, 2);
   TSNode child = ts_node_child(node, (uint32_t)num);
 
-  lua_pushvalue(L, 1);
-  push_node(L, child);
+  push_node(L, child, 1);
   return 1;
 }
 
@@ -600,8 +641,7 @@ static int node_named_child(lua_State *L)
   long num = lua_tointeger(L, 2);
   TSNode child = ts_node_named_child(node, (uint32_t)num);
 
-  lua_pushvalue(L, 1);
-  push_node(L, child);
+  push_node(L, child, 1);
   return 1;
 }
 
@@ -617,8 +657,7 @@ static int node_descendant_for_range(lua_State *L)
                  (uint32_t)lua_tointeger(L, 5) };
   TSNode child = ts_node_descendant_for_point_range(node, start, end);
 
-  lua_pushvalue(L, 1);
-  push_node(L, child);
+  push_node(L, child, 1);
   return 1;
 }
 
@@ -634,8 +673,7 @@ static int node_named_descendant_for_range(lua_State *L)
                  (uint32_t)lua_tointeger(L, 5) };
   TSNode child = ts_node_named_descendant_for_point_range(node, start, end);
 
-  lua_pushvalue(L, 1);
-  push_node(L, child);
+  push_node(L, child, 1);
   return 1;
 }
 
@@ -646,7 +684,254 @@ static int node_parent(lua_State *L)
     return 0;
   }
   TSNode parent = ts_node_parent(node);
-  push_node(L, parent);
+  push_node(L, parent, 1);
+  return 1;
+}
+
+/// assumes the match table being on top of the stack
+static void set_match(lua_State *L, TSQueryMatch *match, int nodeidx)
+{
+  for (int i = 0; i < match->capture_count; i++) {
+    push_node(L, match->captures[i].node, nodeidx);
+    lua_rawseti(L, -2, match->captures[i].index+1);
+  }
+}
+
+static int query_next_match(lua_State *L)
+{
+  TSLua_cursor *ud = lua_touserdata(L, lua_upvalueindex(1));
+  TSQueryCursor *cursor = ud->cursor;
+
+  TSQuery *query = query_check(L, lua_upvalueindex(3));
+  TSQueryMatch match;
+  if (ts_query_cursor_next_match(cursor, &match)) {
+    lua_pushinteger(L, match.pattern_index+1);  // [index]
+    lua_createtable(L, ts_query_capture_count(query), 2);  // [index, match]
+    set_match(L, &match, lua_upvalueindex(2));
+    return 2;
+  }
+  return 0;
+}
+
+
+static int query_next_capture(lua_State *L)
+{
+  TSLua_cursor *ud = lua_touserdata(L, lua_upvalueindex(1));
+  TSQueryCursor *cursor = ud->cursor;
+
+  TSQuery *query = query_check(L, lua_upvalueindex(3));
+
+  if (ud->predicated_match > -1) {
+    lua_getfield(L, lua_upvalueindex(4), "active");
+    bool active = lua_toboolean(L, -1);
+    lua_pop(L, 1);
+    if (!active) {
+      ts_query_cursor_remove_match(cursor, ud->predicated_match);
+    }
+    ud->predicated_match = -1;
+  }
+
+  TSQueryMatch match;
+  uint32_t capture_index;
+  if (ts_query_cursor_next_capture(cursor, &match, &capture_index)) {
+    TSQueryCapture capture = match.captures[capture_index];
+
+    lua_pushinteger(L, capture.index+1);  // [index]
+    push_node(L, capture.node, lua_upvalueindex(2));  // [index, node]
+
+    uint32_t n_pred;
+    ts_query_predicates_for_pattern(query, match.pattern_index, &n_pred);
+    if (n_pred > 0 && capture_index == 0) {
+      lua_pushvalue(L, lua_upvalueindex(4));  // [index, node, match]
+      set_match(L, &match, lua_upvalueindex(2));
+      lua_pushinteger(L, match.pattern_index+1);
+      lua_setfield(L, -2, "pattern");
+
+      if (match.capture_count > 1) {
+        ud->predicated_match = match.id;
+        lua_pushboolean(L, false);
+        lua_setfield(L, -2, "active");
+      }
+      return 3;
+    }
+    return 2;
+  }
+  return 0;
+}
+
+static int node_rawquery(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, &node)) {
+    return 0;
+  }
+  TSQuery *query = query_check(L, 2);
+  // TODO(bfredl): these are expensive allegedly,
+  // use a reuse list later on?
+  TSQueryCursor *cursor = ts_query_cursor_new();
+  ts_query_cursor_exec(cursor, query, node);
+
+  bool captures = lua_toboolean(L, 3);
+
+  if (lua_gettop(L) >= 4) {
+    int start = luaL_checkinteger(L, 4);
+    int end = lua_gettop(L) >= 5 ? luaL_checkinteger(L, 5) : MAXLNUM;
+    ts_query_cursor_set_point_range(cursor,
+                                    (TSPoint){ start, 0 }, (TSPoint){ end, 0 });
+  }
+
+  TSLua_cursor *ud = lua_newuserdata(L, sizeof(*ud));  // [udata]
+  ud->cursor = cursor;
+  ud->predicated_match = -1;
+
+  lua_getfield(L, LUA_REGISTRYINDEX, "treesitter_querycursor");
+  lua_setmetatable(L, -2);  // [udata]
+  lua_pushvalue(L, 1);  // [udata, node]
+
+  // include query separately, as to keep a ref to it for gc
+  lua_pushvalue(L, 2);  // [udata, node, query]
+
+  if (captures) {
+    // placeholder for match state
+    lua_createtable(L, ts_query_capture_count(query), 2);  // [u, n, q, match]
+    lua_pushcclosure(L, query_next_capture, 4);  // [closure]
+  } else {
+    lua_pushcclosure(L, query_next_match, 3);  // [closure]
+  }
+
   return 1;
 }
 
+static int querycursor_gc(lua_State *L)
+{
+  TSLua_cursor *ud = luaL_checkudata(L, 1, "treesitter_querycursor");
+  ts_query_cursor_delete(ud->cursor);
+  return 0;
+}
+
+// Query methods
+
+int ts_lua_parse_query(lua_State *L)
+{
+  if (lua_gettop(L) < 2 || !lua_isstring(L, 1) || !lua_isstring(L, 2)) {
+    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);
+  }
+
+  size_t len;
+  const char *src = lua_tolstring(L, 2, &len);
+
+  uint32_t error_offset;
+  TSQueryError error_type;
+  TSQuery *query = ts_query_new(lang, src, len, &error_offset, &error_type);
+
+  if (!query) {
+    return luaL_error(L, "query: %s at position %d",
+                      query_err_string(error_type), (int)error_offset);
+  }
+
+  TSQuery **ud = lua_newuserdata(L, sizeof(TSQuery *));  // [udata]
+  *ud = query;
+  lua_getfield(L, LUA_REGISTRYINDEX, "treesitter_query");  // [udata, meta]
+  lua_setmetatable(L, -2);  // [udata]
+  return 1;
+}
+
+
+static const char *query_err_string(TSQueryError err) {
+  switch (err) {
+    case TSQueryErrorSyntax: return "invalid syntax";
+    case TSQueryErrorNodeType: return "invalid node type";
+    case TSQueryErrorField: return "invalid field";
+    case TSQueryErrorCapture: return "invalid capture";
+    default: return "error";
+  }
+}
+
+static TSQuery *query_check(lua_State *L, int index)
+{
+  TSQuery **ud = luaL_checkudata(L, index, "treesitter_query");
+  return *ud;
+}
+
+static int query_gc(lua_State *L)
+{
+  TSQuery *query = query_check(L, 1);
+  if (!query) {
+    return 0;
+  }
+
+  ts_query_delete(query);
+  return 0;
+}
+
+static int query_tostring(lua_State *L)
+{
+  lua_pushstring(L, "<query>");
+  return 1;
+}
+
+static int query_inspect(lua_State *L)
+{
+  TSQuery *query = query_check(L, 1);
+  if (!query) {
+    return 0;
+  }
+
+  uint32_t n_pat = ts_query_pattern_count(query);
+  lua_createtable(L, 0, 2);  // [retval]
+  lua_createtable(L, n_pat, 1);  // [retval, patterns]
+  for (size_t i = 0; i < n_pat; i++) {
+    uint32_t len;
+    const TSQueryPredicateStep *step = ts_query_predicates_for_pattern(query,
+                                                                       i, &len);
+    if (len == 0) {
+      continue;
+    }
+    lua_createtable(L, len/4, 1);  // [retval, patterns, pat]
+    lua_createtable(L, 3, 0);  // [retval, patterns, pat, pred]
+    int nextpred = 1;
+    int nextitem = 1;
+    for (size_t k = 0; k < len; k++) {
+      if (step[k].type == TSQueryPredicateStepTypeDone) {
+        lua_rawseti(L, -2, nextpred++);  // [retval, patterns, pat]
+        lua_createtable(L, 3, 0);  // [retval, patterns, pat, pred]
+        nextitem = 1;
+        continue;
+      }
+
+      if (step[k].type == TSQueryPredicateStepTypeString) {
+        uint32_t strlen;
+        const char *str = ts_query_string_value_for_id(query, step[k].value_id,
+                                                       &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]
+      } else {
+        abort();
+      }
+      lua_rawseti(L, -2, nextitem++);  // [retval, patterns, pat, pred]
+    }
+    // last predicate should have ended with TypeDone
+    lua_pop(L, 1);  // [retval, patters, pat]
+    lua_rawseti(L, -2, i+1);  // [retval, patterns]
+  }
+  lua_setfield(L, -2, "patterns");  // [retval]
+
+  uint32_t n_captures = ts_query_capture_count(query);
+  lua_createtable(L, n_captures, 0);  // [retval, captures]
+  for (size_t i = 0; i < n_captures; i++) {
+    uint32_t strlen;
+    const char *str = ts_query_capture_name_for_id(query, i, &strlen);
+    lua_pushlstring(L, str, strlen);  // [retval, captures, capture]
+    lua_rawseti(L, -2, i+1);
+  }
+  lua_setfield(L, -2, "captures");  // [retval]
+
+  return 1;
+}
diff --git a/src/nvim/screen.c b/src/nvim/screen.c
index 1d29ae064e..4082208dd4 100644
--- a/src/nvim/screen.c
+++ b/src/nvim/screen.c
@@ -116,6 +116,8 @@
 #include "nvim/window.h"
 #include "nvim/os/time.h"
 #include "nvim/api/private/helpers.h"
+#include "nvim/api/vim.h"
+#include "nvim/lua/executor.h"
 
 #define MB_FILLER_CHAR '<'  /* character used when a double-width character
                              * doesn't fit. */
@@ -232,6 +234,22 @@ void redraw_buf_line_later(buf_T *buf,  linenr_T line)
   }
 }
 
+void redraw_buf_range_later(buf_T *buf,  linenr_T firstline, linenr_T lastline)
+{
+  FOR_ALL_WINDOWS_IN_TAB(wp, curtab) {
+    if (wp->w_buffer == buf
+        && lastline >= wp->w_topline && firstline < wp->w_botline) {
+      if (wp->w_redraw_top == 0 || wp->w_redraw_top > firstline) {
+          wp->w_redraw_top = firstline;
+      }
+      if (wp->w_redraw_bot == 0 || wp->w_redraw_bot < lastline) {
+          wp->w_redraw_bot = lastline;
+      }
+      redraw_win_later(wp, VALID);
+    }
+  }
+}
+
 /*
  * Changed something in the current window, at buffer line "lnum", that
  * requires that line and possibly other lines to be redrawn.
@@ -477,6 +495,19 @@ int update_screen(int type)
       if (wwp == wp && syntax_present(wp)) {
         syn_stack_apply_changes(wp->w_buffer);
       }
+
+      buf_T *buf = wp->w_buffer;
+      if (buf->b_luahl && buf->b_luahl_window != LUA_NOREF) {
+        Error err = ERROR_INIT;
+        FIXED_TEMP_ARRAY(args, 2);
+        args.items[0] = BUFFER_OBJ(buf->handle);
+        args.items[1] = INTEGER_OBJ(display_tick);
+        executor_exec_lua_cb(buf->b_luahl_start, "start", args, false, &err);
+        if (ERROR_SET(&err)) {
+          ELOG("error in luahl start: %s", err.msg);
+          api_clear_error(&err);
+        }
+      }
     }
   }
 
@@ -1181,7 +1212,27 @@ static void win_update(win_T *wp)
   idx = 0;              /* first entry in w_lines[].wl_size */
   row = 0;
   srow = 0;
-  lnum = wp->w_topline;         /* first line shown in window */
+  lnum = wp->w_topline;  // first line shown in window
+
+  if (buf->b_luahl && buf->b_luahl_window != LUA_NOREF) {
+    Error err = ERROR_INIT;
+    FIXED_TEMP_ARRAY(args, 4);
+    linenr_T knownmax = ((wp->w_valid & VALID_BOTLINE)
+                         ? wp->w_botline
+                         : (wp->w_topline + wp->w_height_inner));
+    args.items[0] = WINDOW_OBJ(wp->handle);
+    args.items[1] = BUFFER_OBJ(buf->handle);
+    args.items[2] = INTEGER_OBJ(wp->w_topline-1);
+    args.items[3] = INTEGER_OBJ(knownmax);
+    // TODO(bfredl): we could allow this callback to change mod_top, mod_bot.
+    // For now the "start" callback is expected to use nvim__buf_redraw_range.
+    executor_exec_lua_cb(buf->b_luahl_window, "window", args, false, &err);
+    if (ERROR_SET(&err)) {
+      ELOG("error in luahl window: %s", err.msg);
+      api_clear_error(&err);
+    }
+  }
+
   for (;; ) {
     /* stop updating when reached the end of the window (check for _past_
      * the end of the window is at the end of the loop) */
@@ -2229,6 +2280,8 @@ win_line (
 
   row = startrow;
 
+  char *luatext = NULL;
+
   if (!number_only) {
     // To speed up the loop below, set extra_check when there is linebreak,
     // trailing white space and/or syntax processing to be done.
@@ -2454,6 +2507,41 @@ win_line (
   line = ml_get_buf(wp->w_buffer, lnum, FALSE);
   ptr = line;
 
+  buf_T *buf = wp->w_buffer;
+  if (buf->b_luahl && buf->b_luahl_line != LUA_NOREF) {
+    size_t size = STRLEN(line);
+    if (lua_attr_bufsize < size) {
+      xfree(lua_attr_buf);
+      lua_attr_buf = xcalloc(size, sizeof(*lua_attr_buf));
+      lua_attr_bufsize = size;
+    } else if (lua_attr_buf) {
+      memset(lua_attr_buf, 0, size * sizeof(*lua_attr_buf));
+    }
+    Error err = ERROR_INIT;
+    // TODO(bfredl): build a macro for the "static array" pattern
+    // in buf_updates_send_changes?
+    FIXED_TEMP_ARRAY(args, 3);
+    args.items[0] = WINDOW_OBJ(wp->handle);
+    args.items[1] = BUFFER_OBJ(buf->handle);
+    args.items[2] = INTEGER_OBJ(lnum-1);
+    lua_attr_active = true;
+    extra_check = true;
+    Object o = executor_exec_lua_cb(buf->b_luahl_line, "line",
+                                    args, true, &err);
+    lua_attr_active = false;
+    if (o.type == kObjectTypeString) {
+      // TODO(bfredl): this is a bit of a hack. A final API should use an
+      // "unified" interface where luahl can add both bufhl and virttext
+      luatext = o.data.string.data;
+      do_virttext = true;
+    } else if (ERROR_SET(&err)) {
+      ELOG("error in luahl line: %s", err.msg);
+      luatext = err.msg;
+      do_virttext = true;
+      api_clear_error(&err);
+    }
+  }
+
   if (has_spell && !number_only) {
     // For checking first word with a capital skip white space.
     if (cap_col == 0) {
@@ -3429,6 +3517,10 @@ win_line (
           }
         }
 
+        if (buf->b_luahl && v > 0 && v < (long)lua_attr_bufsize+1) {
+          char_attr = hl_combine_attr(char_attr, lua_attr_buf[v-1]);
+        }
+
         if (wp->w_buffer->terminal) {
           char_attr = hl_combine_attr(term_attrs[vcol], char_attr);
         }
@@ -3917,8 +4009,14 @@ win_line (
         int rightmost_vcol = 0;
         int i;
 
-        VirtText virt_text = do_virttext ? bufhl_info.line->virt_text
-                                        : (VirtText)KV_INITIAL_VALUE;
+        VirtText virt_text;
+        if (luatext) {
+          virt_text = (VirtText)KV_INITIAL_VALUE;
+          kv_push(virt_text, ((VirtTextChunk){ .text = luatext, .hl_id = 0 }));
+        } else {
+          virt_text = do_virttext ? bufhl_info.line->virt_text
+                                  : (VirtText)KV_INITIAL_VALUE;
+        }
         size_t virt_pos = 0;
         LineState s = LINE_STATE((char_u *)"");
         int virt_attr = 0;
@@ -4319,6 +4417,7 @@ win_line (
   }
 
   xfree(p_extra_free);
+  xfree(luatext);
   return row;
 }