feat: add completion to ':lua'

7 files changed, 420 insertions, 1 deletions

diff --git a/src/nvim/ex_docmd.c b/src/nvim/ex_docmd.c
index ccf7dd0f68..629c7df386 100644
--- a/src/nvim/ex_docmd.c
+++ b/src/nvim/ex_docmd.c
@@ -3680,6 +3680,11 @@ const char * set_one_cmd_context(
     xp->xp_pattern = (char_u *)arg;
     break;
 
+  case CMD_lua:
+    xp->xp_context = EXPAND_LUA;
+    xp->xp_pattern = (char_u *)arg;
+    break;
+
   default:
     break;
   }
@@ -5187,6 +5192,7 @@ static const char *command_complete[] =
 #ifdef HAVE_WORKING_LIBINTL
   [EXPAND_LOCALES] = "locale",
 #endif
+  [EXPAND_LUA] = "lua",
   [EXPAND_MAPCLEAR] = "mapclear",
   [EXPAND_MAPPINGS] = "mapping",
   [EXPAND_MENUS] = "menu",
diff --git a/src/nvim/ex_getln.c b/src/nvim/ex_getln.c
index 2aa66f6a8c..96cab0f110 100644
--- a/src/nvim/ex_getln.c
+++ b/src/nvim/ex_getln.c
@@ -69,6 +69,7 @@
 #include "nvim/lib/kvec.h"
 #include "nvim/api/private/helpers.h"
 #include "nvim/highlight_defs.h"
+#include "nvim/lua/executor.h"
 #include "nvim/viml/parser/parser.h"
 #include "nvim/viml/parser/expressions.h"
 
@@ -5106,6 +5107,10 @@ ExpandFromContext (
   if (xp->xp_context == EXPAND_PACKADD) {
     return ExpandPackAddDir(pat, num_file, file);
   }
+  if (xp->xp_context == EXPAND_LUA) {
+    ILOG("PAT %s", pat);
+    return nlua_expand_pat(pat, num_file, file);
+  }
 
   regmatch.regprog = vim_regcomp(pat, p_magic ? RE_MAGIC : 0);
   if (regmatch.regprog == NULL)
diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 3219c02068..5d75a7dc0f 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -1292,6 +1292,67 @@ static void nlua_add_treesitter(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   lua_setfield(lstate, -2, "_ts_parse_query");
 }
 
+int nlua_expand_pat(char_u *pat, int *num_results, char_u ***results)
+{
+  lua_State *const lstate = nlua_enter();
+  int ret = OK;
+
+  // [ vim ]
+  lua_getglobal(lstate, "vim");
+
+  // [ vim, vim._expand_pat ]
+  lua_getfield(lstate, -1, "_expand_pat");
+  luaL_checktype(lstate, -1, LUA_TFUNCTION);
+
+  // [ vim, vim._log_keystroke, buf ]
+  lua_pushlstring(lstate, (const char *)pat, STRLEN(pat));
+
+  if (lua_pcall(lstate, 1, 1, 0)) {
+    nlua_error(
+        lstate,
+        _("Error executing vim._expand_pat: %.*s"));
+  }
+
+  Error err = ERROR_INIT;
+
+  *num_results = 0;
+  *results = NULL;
+
+  Array completions = nlua_pop_Array(lstate, &err);
+  if (ERROR_SET(&err)) {
+    ret = FAIL;
+    goto cleanup;
+  }
+
+  garray_T result_array;
+  ga_init(&result_array, (int)sizeof(char *), 80);
+  for (size_t i = 0; i < completions.size; i++) {
+    Object v = completions.items[i];
+
+    if (v.type != kObjectTypeString) {
+      ret = FAIL;
+      goto cleanup;
+    }
+
+    GA_APPEND(
+        char_u *,
+        &result_array,
+        vim_strsave((char_u *)v.data.string.data));
+  }
+
+  *results = result_array.ga_data;
+  *num_results = result_array.ga_len;
+
+cleanup:
+  api_free_array(completions);
+
+  if (ret == FAIL) {
+    ga_clear(&result_array);
+  }
+
+  return ret;
+}
+
 static int nlua_regex(lua_State *lstate)
 {
   Error err = ERROR_INIT;
diff --git a/src/nvim/lua/vim.lua b/src/nvim/lua/vim.lua
index 80b311de2c..3e601586ea 100644
--- a/src/nvim/lua/vim.lua
+++ b/src/nvim/lua/vim.lua
@@ -534,4 +534,143 @@ function vim._log_keystroke(char)
   end
 end
 
+--- Generate a list of possible completions for the string.
+--- String starts with ^ and then has the pattern.
+---
+---     1. Can we get it to just return things in the global namespace with that name prefix
+---     2. Can we get it to return things from global namespace even with `print(` in front.
+function vim._expand_pat(pat, env)
+  env = env or _G
+
+  pat = string.sub(pat, 2, #pat)
+
+  if pat == '' then
+    local result = vim.tbl_keys(env)
+    table.sort(result)
+    return result
+  end
+
+  -- TODO: We can handle spaces in [] ONLY.
+  --    We should probably do that at some point, just for cooler completion.
+  -- TODO: We can suggest the variable names to go in []
+  --    This would be difficult as well.
+  --    Probably just need to do a smarter match than just `:match`
+
+  -- Get the last part of the pattern
+  local last_part = pat:match("[%w.:_%[%]'\"]+$")
+  if not last_part then return {} end
+
+  local parts, search_index = vim._expand_pat_get_parts(last_part)
+
+  local match_pat = '^' .. string.sub(last_part, search_index, #last_part)
+  local prefix_match_pat = string.sub(pat, 1, #pat - #match_pat + 1) or ''
+
+  local final_env = env
+  for _, part in ipairs(parts) do
+    if type(final_env) ~= 'table' then
+      return {}
+    end
+
+    -- Normally, we just have a string
+    --  Just attempt to get the string directly from the environment
+    if type(part) == "string" then
+      final_env = rawget(final_env, part)
+    else
+      -- However, sometimes you want to use a variable, and complete on it
+      --    With this, you have the power.
+
+      -- MY_VAR = "api"
+      -- vim[MY_VAR]
+      -- -> _G[MY_VAR] -> "api"
+      local result_key = part[1]
+      if not result_key then
+        return {}
+      end
+
+      local result = rawget(env, result_key)
+
+      if result == nil then
+        return {}
+      end
+
+      final_env = rawget(final_env, result)
+    end
+
+    if not final_env then
+      return {}
+    end
+  end
+
+  local result = vim.tbl_map(function(v)
+    return prefix_match_pat .. v
+  end, vim.tbl_filter(function(name)
+    return string.find(name, match_pat) ~= nil
+  end, vim.tbl_keys(final_env)))
+
+  table.sort(result)
+
+  return result
+end
+
+vim._expand_pat_get_parts = function(lua_string)
+  local parts = {}
+
+  local accumulator, search_index = '', 1
+  local in_brackets, bracket_end = false, -1
+  local string_char = nil
+  for idx = 1, #lua_string do
+    local s = lua_string:sub(idx, idx)
+
+    if not in_brackets and (s == "." or s == ":") then
+      table.insert(parts, accumulator)
+      accumulator = ''
+
+      search_index = idx + 1
+    elseif s == "[" then
+      in_brackets = true
+
+      table.insert(parts, accumulator)
+      accumulator = ''
+
+      search_index = idx + 1
+    elseif in_brackets then
+      if idx == bracket_end then
+        in_brackets = false
+        search_index = idx + 1
+
+        if string_char == "VAR" then
+          table.insert(parts, { accumulator })
+          accumulator = ''
+
+          string_char = nil
+        end
+      elseif not string_char then
+        bracket_end = string.find(lua_string, ']', idx, true)
+
+        if s == '"' or s == "'" then
+          string_char = s
+        elseif s ~= ' ' then
+          string_char = "VAR"
+          accumulator = s
+        end
+      elseif string_char then
+        if string_char ~= s then
+          accumulator = accumulator .. s
+        else
+          table.insert(parts, accumulator)
+          accumulator = ''
+
+          string_char = nil
+        end
+      end
+    else
+      accumulator = accumulator .. s
+    end
+  end
+
+  parts = vim.tbl_filter(function(val) return #val > 0 end, parts)
+
+  return parts, search_index
+end
+
 return module
diff --git a/src/nvim/vim.h b/src/nvim/vim.h
index 01f20cf29a..e70749795b 100644
--- a/src/nvim/vim.h
+++ b/src/nvim/vim.h
@@ -159,6 +159,7 @@ enum {
   EXPAND_MAPCLEAR,
   EXPAND_ARGLIST,
   EXPAND_CHECKHEALTH,
+  EXPAND_LUA,
 };
 
 
diff --git a/test/functional/lua/command_line_completion_spec.lua b/test/functional/lua/command_line_completion_spec.lua
new file mode 100644
index 0000000000..9056f80fb5
--- /dev/null
+++ b/test/functional/lua/command_line_completion_spec.lua
@@ -0,0 +1,167 @@
+local helpers = require('test.functional.helpers')(after_each)
+
+local clear = helpers.clear
+local eq = helpers.eq
+local funcs = helpers.funcs
+local exec_lua = helpers.exec_lua
+
+local get_completions = function(input, env)
+  return exec_lua("return vim._expand_pat(...)", '^' .. input, env)
+end
+
+local get_compl_parts = function(parts)
+  return funcs.luaeval("{vim._expand_pat_get_parts(_A)}", parts)
+end
+
+before_each(clear)
+
+describe('nlua_expand_pat', function()
+  it('should complete exact matches', function()
+    eq({'exact'}, get_completions('exact', { exact = true }))
+  end)
+
+  it('should return empty table when nothing matches', function()
+    eq({}, get_completions('foo', { bar = true }))
+  end)
+
+  it('should return nice completions with function call prefix', function()
+    eq({'print(FOO'}, get_completions('print(F', { FOO = true, bawr = true }))
+  end)
+
+  it('should return keys for nested dictionaries', function()
+    eq(
+      {
+        'vim.api.nvim_buf_set_lines',
+        'vim.api.nvim_buf_set_option'
+      },
+      get_completions('vim.api.nvim_buf_', {
+        vim = {
+          api = {
+            nvim_buf_set_lines = true,
+            nvim_buf_set_option = true,
+            nvim_win_doesnt_match = true,
+          },
+          other_key = true,
+        }
+      })
+    )
+  end)
+
+  it('it should work with colons', function()
+    eq(
+      {
+        'MyClass:bawr',
+        'MyClass:baz',
+      },
+      get_completions('MyClass:b', {
+        MyClass = {
+          baz = true,
+          bawr = true,
+          foo = false,
+        }
+      })
+    )
+  end)
+
+  it('should return keys for string reffed dictionaries', function()
+    eq(
+      {
+        'vim["api"].nvim_buf_set_lines',
+        'vim["api"].nvim_buf_set_option'
+      },
+      get_completions('vim["api"].nvim_buf_', {
+        vim = {
+          api = {
+            nvim_buf_set_lines = true,
+            nvim_buf_set_option = true,
+            nvim_win_doesnt_match = true,
+          },
+          other_key = true,
+        }
+      })
+    )
+  end)
+
+  it('should return keys for string reffed dictionaries', function()
+    eq(
+      {
+        'vim["nested"]["api"].nvim_buf_set_lines',
+        'vim["nested"]["api"].nvim_buf_set_option'
+      },
+      get_completions('vim["nested"]["api"].nvim_buf_', {
+        vim = {
+          nested = {
+            api = {
+              nvim_buf_set_lines = true,
+              nvim_buf_set_option = true,
+              nvim_win_doesnt_match = true,
+            },
+          },
+          other_key = true,
+        }
+      })
+    )
+  end)
+
+  it('should be able to interpolate globals', function()
+    eq(
+      {
+        'vim[MY_VAR].nvim_buf_set_lines',
+        'vim[MY_VAR].nvim_buf_set_option'
+      },
+      get_completions('vim[MY_VAR].nvim_buf_', {
+        MY_VAR = "api",
+        vim = {
+          api = {
+            nvim_buf_set_lines = true,
+            nvim_buf_set_option = true,
+            nvim_win_doesnt_match = true,
+          },
+          other_key = true,
+        }
+      })
+    )
+  end)
+
+  it('should return everything if the input is of length 0', function()
+    eq({"other", "vim"}, get_completions('', { vim = true, other = true }))
+  end)
+
+  describe('get_parts', function()
+    it('should return an empty list for no separators', function()
+      eq({{}, 1}, get_compl_parts("vim"))
+    end)
+
+    it('just the first item before a period', function()
+      eq({{"vim"}, 5}, get_compl_parts("vim.ap"))
+    end)
+
+    it('should return multiple parts just for period', function()
+      eq({{"vim", "api"}, 9}, get_compl_parts("vim.api.nvim_buf"))
+    end)
+
+    it('should be OK with colons', function()
+      eq({{"vim", "api"}, 9}, get_compl_parts("vim:api.nvim_buf"))
+    end)
+
+    it('should work for just one string ref', function()
+      eq({{"vim", "api"}, 12}, get_compl_parts("vim['api'].nvim_buf"))
+    end)
+
+    it('should work for just one string ref, with double quote', function()
+      eq({{"vim", "api"}, 12}, get_compl_parts('vim["api"].nvim_buf'))
+    end)
+
+    it('should allows back-to-back string ref', function()
+      eq({{"vim", "nested", "api"}, 22}, get_compl_parts('vim["nested"]["api"].nvim_buf'))
+    end)
+
+    it('should allows back-to-back string ref with spaces before and after', function()
+      eq({{"vim", "nested", "api"}, 25}, get_compl_parts('vim[ "nested"  ]["api"].nvim_buf'))
+    end)
+
+    it('should allow VAR style loolup', function()
+      eq({{"vim", {"NESTED"}, "api"}, 20}, get_compl_parts('vim[NESTED]["api"].nvim_buf'))
+    end)
+  end)
+end)
diff --git a/test/functional/viml/completion_spec.lua b/test/functional/viml/completion_spec.lua
index 01fc50289d..d59a1a3306 100644
--- a/test/functional/viml/completion_spec.lua
+++ b/test/functional/viml/completion_spec.lua
@@ -26,6 +26,7 @@ describe('completion', function()
       [7] = {foreground = Screen.colors.White, background = Screen.colors.Red},
       [8] = {reverse = true},
       [9] = {bold = true, reverse = true},
+      [10] = {foreground = Screen.colors.Grey0, background = Screen.colors.Yellow},
     })
   end)
 
@@ -895,8 +896,47 @@ describe('completion', function()
     ]])
   end)
 
-  describe('from the commandline window', function()
+  describe('lua completion', function()
+    it('expands when there is only one match', function()
+      feed(':lua CURRENT_TESTING_VAR = 1<CR>')
+      feed(':lua CURRENT_TESTING_<TAB>')
+      screen:expect{grid=[[
+                                                                    |
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {0:~                                                           }|
+        :lua CURRENT_TESTING_VAR^                                    |
+      ]]}
+    end)
 
+    it('expands when there is only one match', function()
+      feed(':lua CURRENT_TESTING_FOO = 1<CR>')
+      feed(':lua CURRENT_TESTING_BAR = 1<CR>')
+      feed(':lua CURRENT_TESTING_<TAB>')
+      screen:expect{ grid = [[
+                                                                    |
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {0:~                                                           }|
+        {10:CURRENT_TESTING_BAR}{9:  CURRENT_TESTING_FOO                    }|
+        :lua CURRENT_TESTING_BAR^                                    |
+      ]], unchanged = true }
+    end)
+
+    it('provides completion from `getcompletion()`', function()
+      eq({'vim'}, meths.call_function('getcompletion', {'vi', 'lua'}))
+      eq({'vim.api'}, meths.call_function('getcompletion', {'vim.ap', 'lua'}))
+      eq({'vim.tbl_filter'}, meths.call_function('getcompletion', {'vim.tbl_fil', 'lua'}))
+      eq({'print(vim'}, meths.call_function('getcompletion', {'print(vi', 'lua'}))
+    end)
+  end)
+
+  describe('from the commandline window', function()
     it('is cleared after CTRL-C', function ()
       feed('q:')
       feed('ifoo faa fee f')