lua: add regex support, and `@match` support in treesitter queries

6 files changed, 252 insertions, 17 deletions

diff --git a/runtime/doc/lua.txt b/runtime/doc/lua.txt
index c113a70027..800f24b5c9 100644
--- a/runtime/doc/lua.txt
+++ b/runtime/doc/lua.txt
@@ -693,6 +693,35 @@ identical identifiers, highlighting both as |hl-WarningMsg|: >
      (eq? @WarningMsg.left @WarningMsg.right))
 
 ------------------------------------------------------------------------------
+VIM.REGEX							*lua-regex*
+
+Vim regexes can be used directly from lua. Currently they only allow
+matching within a single line.
+
+vim.regex({re})						*vim.regex()*
+
+        Parse the regex {re} and return a regex object. 'magic' and
+        'ignorecase' options are ignored, lua regexes always defaults to magic
+        and ignoring case.  The behavior can be changed with flags in
+        the beginning of the string |/magic|.
+
+Regex objects support the following methods:
+
+regex:match_str({str})					*regex:match_str()*
+        Match the string against the regex. If the string should match the
+        regex precisely, surround the regex with `^` and `$`.
+        If the was a match, the byte indices for the beginning and end of
+        the match is returned. When there is no match, `nil` is returned.
+        As any integer is truth-y, `regex:match()` can be directly used
+        as a condition in an if-statement.
+
+regex:match_line({bufnr}, {line_idx}[, {start}, {end}])	*regex:match_line()*
+        Match line {line_idx} (zero-based) in buffer {bufnr}. If {start} and
+        {end} are supplied, match only this byte index range. Otherwise see
+        |regex:match_str()|. If {start} is used, then the returned byte
+        indices will be relative {start}.
+
+------------------------------------------------------------------------------
 VIM							*lua-builtin*
 
 vim.api.{func}({...})					*vim.api*
diff --git a/runtime/lua/vim/treesitter.lua b/runtime/lua/vim/treesitter.lua
index 0d0e22adb3..8dacfa11cf 100644
--- a/runtime/lua/vim/treesitter.lua
+++ b/runtime/lua/vim/treesitter.lua
@@ -113,12 +113,33 @@ end
 local Query = {}
 Query.__index = Query
 
+local magic_prefixes = {['\\v']=true, ['\\m']=true, ['\\M']=true, ['\\V']=true}
+local function check_magic(str)
+  if string.len(str) < 2 or magic_prefixes[string.sub(str,1,2)] then
+    return str
+  end
+  return '\\v'..str
+end
+
 function M.parse_query(lang, query)
   M.require_language(lang)
   local self = setmetatable({}, Query)
   self.query = vim._ts_parse_query(lang, query)
   self.info = self.query:inspect()
   self.captures = self.info.captures
+  self.regexes = {}
+  for id,preds in pairs(self.info.patterns) do
+    local regexes = {}
+    for i, pred in ipairs(preds) do
+      if (pred[1] == "match?" and type(pred[2]) == "number"
+          and type(pred[3]) == "string") then
+        regexes[i] = vim.regex(check_magic(pred[3]))
+      end
+    end
+    if next(regexes) then
+      self.regexes[id] = regexes
+    end
+  end
   return self
 end
 
@@ -131,8 +152,13 @@ local function get_node_text(node, bufnr)
   return string.sub(line, start_col+1, end_col)
 end
 
-local function match_preds(match, preds, bufnr)
-  for _, pred in pairs(preds) do
+function Query:match_preds(match, pattern, bufnr)
+  local preds = self.info.patterns[pattern]
+  if not preds then
+    return true
+  end
+  local regexes = self.regexes[pattern]
+  for i, pred in pairs(preds) do
     if pred[1] == "eq?" then
       local node = match[pred[2]]
       local node_text = get_node_text(node, bufnr)
@@ -149,6 +175,16 @@ local function match_preds(match, preds, bufnr)
       if node_text ~= str or str == nil then
         return false
       end
+    elseif pred[1] == "match?" then
+      if not regexes or not regexes[i] then
+        return false
+      end
+      local node = match[pred[2]]
+      local start_row, start_col, end_row, end_col = node:range()
+      if start_row ~= end_row then
+        return false
+      end
+      return regexes[i]:match_line(bufnr, start_row, start_col, end_col)
     else
       return false
     end
@@ -164,8 +200,7 @@ function Query:iter_captures(node, bufnr, start, stop)
   local function iter()
     local capture, captured_node, match = raw_iter()
     if match ~= nil then
-      local preds = self.info.patterns[match.pattern]
-      local active = match_preds(match, preds, bufnr)
+      local active = self:match_preds(match, match.pattern, bufnr)
       match.active = active
       if not active then
         return iter() -- tail call: try next match
@@ -184,8 +219,7 @@ function Query:iter_matches(node, bufnr, start, stop)
   local function iter()
     local pattern, match = raw_iter()
     if match ~= nil then
-      local preds = self.info.patterns[pattern]
-      local active = (not preds) or match_preds(match, preds, bufnr)
+      local active = self:match_preds(match, pattern, bufnr)
       if not active then
         return iter() -- tail call: try next match
       end
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 9e0063ebaa..9167ec5e54 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -11,6 +11,7 @@
 #include "nvim/func_attr.h"
 #include "nvim/api/private/defs.h"
 #include "nvim/api/private/helpers.h"
+#include "nvim/api/private/handle.h"
 #include "nvim/api/vim.h"
 #include "nvim/msgpack_rpc/channel.h"
 #include "nvim/vim.h"
@@ -19,6 +20,7 @@
 #include "nvim/message.h"
 #include "nvim/memline.h"
 #include "nvim/buffer_defs.h"
+#include "nvim/regexp.h"
 #include "nvim/macros.h"
 #include "nvim/screen.h"
 #include "nvim/cursor.h"
@@ -244,6 +246,14 @@ static int nlua_schedule(lua_State *const lstate)
   return 0;
 }
 
+static struct luaL_Reg regex_meta[] = {
+  { "__gc", regex_gc },
+  { "__tostring", regex_tostring },
+  { "match_str", regex_match_str },
+  { "match_line", regex_match_line },
+  { NULL, NULL }
+};
+
 /// Initialize lua interpreter state
 ///
 /// Called by lua interpreter itself to initialize state.
@@ -291,6 +301,15 @@ static int nlua_state_init(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   // call
   lua_pushcfunction(lstate, &nlua_call);
   lua_setfield(lstate, -2, "call");
+  // regex
+  lua_pushcfunction(lstate, &nlua_regex);
+  lua_setfield(lstate, -2, "regex");
+
+  luaL_newmetatable(lstate, "nvim_regex");
+  luaL_register(lstate, NULL, regex_meta);
+  lua_pushvalue(lstate, -1);  // [meta, meta]
+  lua_setfield(lstate, -2, "__index");  // [meta]
+  lua_pop(lstate, 1);  // don't use metatable now
 
   // rpcrequest
   lua_pushcfunction(lstate, &nlua_rpcrequest);
@@ -1037,3 +1056,136 @@ static void nlua_add_treesitter(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   lua_pushcfunction(lstate, ts_lua_parse_query);
   lua_setfield(lstate, -2, "_ts_parse_query");
 }
+
+static int nlua_regex(lua_State *lstate)
+{
+  Error err = ERROR_INIT;
+  const char *text = luaL_checkstring(lstate, 1);
+  regprog_T *prog = NULL;
+
+  TRY_WRAP({
+    try_start();
+    prog = vim_regcomp((char_u *)text, RE_AUTO | RE_MAGIC | RE_STRICT);
+    try_end(&err);
+  });
+
+  if (ERROR_SET(&err)) {
+    return luaL_error(lstate, "couldn't parse regex: %s", err.msg);
+  }
+  assert(prog);
+
+  regprog_T **p = lua_newuserdata(lstate, sizeof(regprog_T *));
+  *p = prog;
+
+  lua_getfield(lstate, LUA_REGISTRYINDEX, "nvim_regex");  // [udata, meta]
+  lua_setmetatable(lstate, -2);  // [udata]
+  return 1;
+}
+
+static regprog_T **regex_check(lua_State *L)
+{
+  return luaL_checkudata(L, 1, "nvim_regex");
+}
+
+
+static int regex_gc(lua_State *lstate)
+{
+  regprog_T **prog = regex_check(lstate);
+  vim_regfree(*prog);
+  return 0;
+}
+
+static int regex_tostring(lua_State *lstate)
+{
+  lua_pushstring(lstate, "<regex>");
+  return 1;
+}
+
+static int regex_match(lua_State *lstate, regprog_T **prog, char_u *str)
+{
+  regmatch_T rm;
+  rm.regprog = *prog;
+  rm.rm_ic = false;
+  bool match = vim_regexec(&rm, str, 0);
+  *prog = rm.regprog;
+
+  if (match) {
+    lua_pushinteger(lstate, (lua_Integer)(rm.startp[0]-str));
+    lua_pushinteger(lstate, (lua_Integer)(rm.endp[0]-str));
+    return 2;
+  }
+  return 0;
+}
+
+static int regex_match_str(lua_State *lstate)
+{
+  regprog_T **prog = regex_check(lstate);
+  const char *str = luaL_checkstring(lstate, 2);
+  int nret = regex_match(lstate, prog, (char_u *)str);
+
+  if (!*prog) {
+    return luaL_error(lstate, "regex: internal error");
+  }
+
+  return nret;
+}
+
+static int regex_match_line(lua_State *lstate)
+{
+  regprog_T **prog = regex_check(lstate);
+
+  int narg = lua_gettop(lstate);
+  if (narg < 3) {
+    return luaL_error(lstate, "not enough args");
+  }
+
+  long bufnr = luaL_checkinteger(lstate, 2);
+  long rownr = luaL_checkinteger(lstate, 3);
+  long start = 0, end = -1;
+  if (narg >= 4) {
+    start = luaL_checkinteger(lstate, 4);
+  }
+  if (narg >= 5) {
+    end = luaL_checkinteger(lstate, 5);
+    if (end < 0) {
+      return luaL_error(lstate, "invalid end");
+    }
+  }
+
+  buf_T *buf = bufnr ? handle_get_buffer((int)bufnr) : curbuf;
+  if (!buf || buf->b_ml.ml_mfp == NULL) {
+    return luaL_error(lstate, "invalid buffer");
+  }
+
+  if (rownr >= buf->b_ml.ml_line_count) {
+    return luaL_error(lstate, "invalid row");
+  }
+
+  char_u *line = ml_get_buf(buf, rownr+1, false);
+  size_t len = STRLEN(line);
+
+  if (start < 0 || (size_t)start > len) {
+    return luaL_error(lstate, "invalid start");
+  }
+
+  char_u save = NUL;
+  if (end >= 0) {
+    if ((size_t)end > len || end < start) {
+      return luaL_error(lstate, "invalid end");
+    }
+    save = line[end];
+    line[end] = NUL;
+  }
+
+  int nret = regex_match(lstate, prog, line+start);
+
+  if (end >= 0) {
+    line[end] = save;
+  }
+
+  if (!*prog) {
+    return luaL_error(lstate, "regex: internal error");
+  }
+
+  return nret;
+}
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index a420f79ffd..4753df7b87 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -93,10 +93,7 @@ static PMap(cstr_t) *langs;
 static void build_meta(lua_State *L, const char *tname, const luaL_Reg *meta)
 {
   if (luaL_newmetatable(L, tname)) {  // [meta]
-    for (size_t i = 0; meta[i].name != NULL; i++) {
-      lua_pushcfunction(L, meta[i].func);  // [meta, func]
-      lua_setfield(L, -2, meta[i].name);  // [meta]
-    }
+    luaL_register(L, NULL, meta);
 
     lua_pushvalue(L, -1);  // [meta, meta]
     lua_setfield(L, -2, "__index");  // [meta]
diff --git a/test/functional/lua/treesitter_spec.lua b/test/functional/lua/treesitter_spec.lua
index 494d6c84bb..f93185d1f6 100644
--- a/test/functional/lua/treesitter_spec.lua
+++ b/test/functional/lua/treesitter_spec.lua
@@ -245,6 +245,11 @@ static int nlua_schedule(lua_State *const lstate)
 (primitive_type) @type
 (sized_type_specifier) @type
 
+; defaults to very magic syntax, for best compatibility
+((identifier) @Identifier (match? @Identifier "^l(u)a_"))
+; still support \M etc prefixes
+((identifier) @Constant (match? @Constant "\M^\[A-Z_]\+$"))
+
 ((binary_expression left: (identifier) @WarningMsg.left right: (identifier) @WarningMsg.right) (eq? @WarningMsg.left @WarningMsg.right))
 
 (comment) @comment
@@ -263,7 +268,7 @@ static int nlua_schedule(lua_State *const lstate)
       [8] = {foreground = Screen.colors.Grey100, background = Screen.colors.Red},
       [9] = {foreground = Screen.colors.Magenta, background = Screen.colors.Red},
       [10] = {foreground = Screen.colors.Red, background = Screen.colors.Red},
-
+      [11] = {foreground = Screen.colors.Cyan4},
     })
 
     insert(hl_text)
@@ -297,10 +302,10 @@ static int nlua_schedule(lua_State *const lstate)
       {2:/// Schedule Lua callback on main loop's event queue}             |
       {3:static} {3:int} nlua_schedule(lua_State *{3:const} lstate)                |
       {                                                                |
-        {4:if} (lua_type(lstate, {5:1}) != LUA_TFUNCTION                       |
+        {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
             || {6:lstate} != {6:lstate}) {                                     |
-          lua_pushliteral(lstate, {5:"vim.schedule: expected function"});  |
-          {4:return} lua_error(lstate);                                    |
+          {11:lua_pushliteral}(lstate, {5:"vim.schedule: expected function"});  |
+          {4:return} {11:lua_error}(lstate);                                    |
         }                                                              |
                                                                        |
         {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
@@ -319,10 +324,10 @@ static int nlua_schedule(lua_State *const lstate)
       {2:/// Schedule Lua callback on main loop's event queue}             |
       {3:static} {3:int} nlua_schedule(lua_State *{3:const} lstate)                |
       {                                                                |
-        {4:if} (lua_type(lstate, {5:1}) != LUA_TFUNCTION                       |
+        {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
             || {6:lstate} != {6:lstate}) {                                     |
-          lua_pushliteral(lstate, {5:"vim.schedule: expected function"});  |
-          {4:return} lua_error(lstate);                                    |
+          {11:lua_pushliteral}(lstate, {5:"vim.schedule: expected function"});  |
+          {4:return} {11:lua_error}(lstate);                                    |
       {8:*^/}                                                               |
         }                                                              |
                                                                        |
diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua
index 2f459145f5..532368c37c 100644
--- a/test/functional/lua/vim_spec.lua
+++ b/test/functional/lua/vim_spec.lua
@@ -844,4 +844,22 @@ describe('lua stdlib', function()
     eq('2', funcs.luaeval "BUF")
     eq(2, funcs.luaeval "#vim.api.nvim_list_bufs()")
   end)
+
+  it('vim.regex', function()
+    exec_lua [[
+      re1 = vim.regex"ab\\+c"
+      vim.cmd "set nomagic ignorecase"
+      re2 = vim.regex"xYz"
+    ]]
+    eq({}, exec_lua[[return {re1:match_str("x ac")}]])
+    eq({3,7}, exec_lua[[return {re1:match_str("ac abbc")}]])
+
+    meths.buf_set_lines(0, 0, -1, true, {"yy", "abc abbc"})
+    eq({}, exec_lua[[return {re1:match_line(0, 0)}]])
+    eq({0,3}, exec_lua[[return {re1:match_line(0, 1)}]])
+    eq({3,7}, exec_lua[[return {re1:match_line(0, 1, 1)}]])
+    eq({3,7}, exec_lua[[return {re1:match_line(0, 1, 1, 8)}]])
+    eq({}, exec_lua[[return {re1:match_line(0, 1, 1, 7)}]])
+    eq({0,3}, exec_lua[[return {re1:match_line(0, 1, 0, 7)}]])
+  end)
 end)