feat(lua): vim.ui_attach to get ui events from lua

Co-authored-by: Famiu Haque <famiuhaque@protonmail.com>

runtime/doc/lua.txt

@@ -967,6 +967,37 @@ vim.wait({time} [, {callback}, {interval}, {fast_only}])          *vim.wait()*
     end
 <
 
+vim.ui_attach({ns}, {options}, {callback})                *vim.ui_attach()*
+    Attach to ui events, similar to |nvim_ui_attach()| but receive events
+    as lua callback. Can be used to implement screen elements like
+    popupmenu or message handling in lua.
+
+    {options} should be a dictionary-like table, where `ext_...` options should
+    be set to true to receive events for the respective external element.
+
+    {callback} receives event name plus additional parameters. See |ui-popupmenu|
+    and the sections below for event format for respective events.
+
+    Example (stub for a |ui-popupmenu| implementation): >
+
+      ns = vim.api.nvim_create_namespace('my_fancy_pum')
+
+      vim.ui_attach(ns, {ext_popupmenu=true}, function(event, ...)
+        if event == "popupmenu_show" then
+          local items, selected, row, col, grid = ...
+          print("display pum ", #items)
+        elseif event == "popupmenu_select" then
+          local selected = ...
+          print("selected", selected)
+        elseif event == "popupmenu_hide" then
+          print("FIN")
+        end
+      end)
+
+vim.ui_detach({ns})                                           *vim.ui_detach()*
+    Detach a callback previously attached with |vim.ui_attach()| for the
+    given namespace {ns}.
+
 vim.type_idx                                                    *vim.type_idx*
     Type index for use in |lua-special-tbl|. Specifying one of the values from
     |vim.types| allows typing the empty table (it is unclear whether empty Lua

src/nvim/api/extmark.c

@@ -87,7 +87,7 @@ const char *describe_ns(NS ns_id)
 }
 
 // Is the Namespace in use?
-static bool ns_initialized(uint32_t ns)
+bool ns_initialized(uint32_t ns)
 {
   if (ns < 1) {
     return false;

src/nvim/api/ui.c

@@ -223,6 +223,7 @@ void nvim_ui_attach(uint64_t channel_id, Integer width, Integer height, Dictiona
   ui->msg_set_pos = remote_ui_msg_set_pos;
   ui->event = remote_ui_event;
   ui->inspect = remote_ui_inspect;
+  ui->win_viewport = remote_ui_win_viewport;
 
   CLEAR_FIELD(ui->ui_ext);
 

src/nvim/api/ui_events.in.h

@@ -100,7 +100,7 @@ void raw_line(Integer grid, Integer row, Integer startcol,
   FUNC_API_NOEXPORT FUNC_API_COMPOSITOR_IMPL;
 
 void event(char *name, Array args)
-  FUNC_API_NOEXPORT;
+  FUNC_API_NOEXPORT FUNC_API_COMPOSITOR_IMPL;
 
 void win_pos(Integer grid, Window win, Integer startrow,
              Integer startcol, Integer width, Integer height)
@@ -121,7 +121,7 @@ void msg_set_pos(Integer grid, Integer row, Boolean scrolled, String sep_char)
 void win_viewport(Integer grid, Window win, Integer topline,
                   Integer botline, Integer curline, Integer curcol,
                   Integer line_count)
-  FUNC_API_SINCE(7) FUNC_API_REMOTE_ONLY;
+  FUNC_API_SINCE(7) FUNC_API_BRIDGE_IMPL;
 
 void win_extmark(Integer grid, Window win, Integer ns_id, Integer mark_id,
                  Integer row, Integer col)

src/nvim/buffer_updates.c

@@ -285,14 +285,13 @@ void buf_updates_send_changes(buf_T *buf, linenr_T firstline, int64_t num_added,
         args.items[7] = INTEGER_OBJ((Integer)deleted_codeunits);
       }
       textlock++;
-      Object res = nlua_call_ref(cb.on_lines, "lines", args, true, NULL);
+      Object res = nlua_call_ref(cb.on_lines, "lines", args, false, NULL);
       textlock--;
 
       if (res.type == kObjectTypeBoolean && res.data.boolean == true) {
         buffer_update_callbacks_free(cb);
         keep = false;
       }
-      api_free_object(res);
     }
     if (keep) {
       kv_A(buf->update_callbacks, j++) = kv_A(buf->update_callbacks, i);
@@ -335,7 +334,7 @@ void buf_updates_send_splice(buf_T *buf, int start_row, colnr_T start_col, bcoun
       ADD_C(args, INTEGER_OBJ(new_byte));
 
       textlock++;
-      Object res = nlua_call_ref(cb.on_bytes, "bytes", args, true, NULL);
+      Object res = nlua_call_ref(cb.on_bytes, "bytes", args, false, NULL);
       textlock--;
 
       if (res.type == kObjectTypeBoolean && res.data.boolean == true) {
@@ -371,14 +370,13 @@ void buf_updates_changedtick(buf_T *buf)
 
       textlock++;
       Object res = nlua_call_ref(cb.on_changedtick, "changedtick",
-                                 args, true, NULL);
+                                 args, false, NULL);
       textlock--;
 
       if (res.type == kObjectTypeBoolean && res.data.boolean == true) {
         buffer_update_callbacks_free(cb);
         keep = false;
       }
-      api_free_object(res);
     }
     if (keep) {
       kv_A(buf->update_callbacks, j++) = kv_A(buf->update_callbacks, i);

src/nvim/eval.c

@@ -5858,13 +5858,8 @@ bool callback_call(Callback *const callback, const int argcount_in, typval_T *co
     break;
 
   case kCallbackLua:
-    rv = nlua_call_ref(callback->data.luaref, NULL, args, true, NULL);
-    switch (rv.type) {
-    case kObjectTypeBoolean:
-      return rv.data.boolean;
-    default:
-      return false;
-    }
+    rv = nlua_call_ref(callback->data.luaref, NULL, args, false, NULL);
+    return (rv.type == kObjectTypeBoolean && rv.data.boolean == true);
 
   case kCallbackNone:
     return false;

src/nvim/generators/gen_api_ui_events.lua

@@ -75,6 +75,8 @@ local function call_ui_event_method(output, ev)
       hlattrs_args_count = hlattrs_args_count + 1
     elseif kind == 'Object' then
       output:write('args.items['..(j-1)..'];\n')
+    elseif kind == 'Window' then
+      output:write('(Window)args.items['..(j-1)..'].data.integer;\n')
     else
       output:write('args.items['..(j-1)..'].data.'..string.lower(kind)..';\n')
     end

src/nvim/lua/executor.c

@@ -7,6 +7,7 @@
 #include <tree_sitter/api.h>
 
 #include "luv/luv.h"
+#include "nvim/api/extmark.h"
 #include "nvim/api/private/defs.h"
 #include "nvim/api/private/helpers.h"
 #include "nvim/api/vim.h"
@@ -40,6 +41,9 @@
 #include "nvim/os/os.h"
 #include "nvim/profile.h"
 #include "nvim/runtime.h"
+#include "nvim/screen.h"
+#include "nvim/ui.h"
+#include "nvim/ui_compositor.h"
 #include "nvim/undo.h"
 #include "nvim/usercmd.h"
 #include "nvim/version.h"
@@ -589,6 +593,71 @@ static bool nlua_init_packages(lua_State *lstate)
   return true;
 }
 
+/// "vim.ui_attach(ns_id, {ext_foo=true}, cb)" function
+static int nlua_ui_attach(lua_State *lstate)
+  FUNC_ATTR_NONNULL_ALL
+{
+  uint32_t ns_id = (uint32_t)luaL_checkinteger(lstate, 1);
+
+  if (!ns_initialized(ns_id)) {
+    return luaL_error(lstate, "invalid ns_id");
+  }
+  if (!lua_istable(lstate, 2)) {
+    return luaL_error(lstate, "ext_widgets must be a table");
+  }
+  if (!lua_isfunction(lstate, 3)) {
+    return luaL_error(lstate, "callback must be a Lua function");
+  }
+
+  bool ext_widgets[kUIGlobalCount] = { false };
+  bool tbl_has_true_val = false;
+
+  lua_pushvalue(lstate, 2);
+  lua_pushnil(lstate);
+  while (lua_next(lstate, -2)) {
+    // [dict, key, val]
+    size_t len;
+    const char *s = lua_tolstring(lstate, -2, &len);
+    bool val = lua_toboolean(lstate, -1);
+
+    for (size_t i = 0; i < kUIGlobalCount; i++) {
+      if (strequal(s, ui_ext_names[i])) {
+        if (val) {
+          tbl_has_true_val = true;
+        }
+        ext_widgets[i] = val;
+        goto ok;
+      }
+    }
+
+    return luaL_error(lstate, "Unexpected key: %s", s);
+ok:
+    lua_pop(lstate, 1);
+  }
+
+  if (!tbl_has_true_val) {
+    return luaL_error(lstate, "ext_widgets table must contain at least one 'true' value");
+  }
+
+  LuaRef ui_event_cb = nlua_ref_global(lstate, 3);
+  ui_comp_add_cb(ns_id, ui_event_cb, ext_widgets);
+  return 0;
+}
+
+/// "vim.ui_detach(ns_id)" function
+static int nlua_ui_detach(lua_State *lstate)
+  FUNC_ATTR_NONNULL_ALL
+{
+  uint32_t ns_id = (uint32_t)luaL_checkinteger(lstate, 1);
+
+  if (!ns_initialized(ns_id)) {
+    return luaL_error(lstate, "invalid ns_id");
+  }
+
+  ui_comp_remove_cb(ns_id);
+  return 0;
+}
+
 /// Initialize lua interpreter state
 ///
 /// Called by lua interpreter itself to initialize state.
@@ -649,6 +718,14 @@ static bool nlua_state_init(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   lua_pushcfunction(lstate, &nlua_wait);
   lua_setfield(lstate, -2, "wait");
 
+  // ui_attach
+  lua_pushcfunction(lstate, &nlua_ui_attach);
+  lua_setfield(lstate, -2, "ui_attach");
+
+  // ui_detach
+  lua_pushcfunction(lstate, &nlua_ui_detach);
+  lua_setfield(lstate, -2, "ui_detach");
+
   nlua_common_vim_init(lstate, false);
 
   // patch require() (only for --startuptime)
@@ -1422,9 +1499,10 @@ bool nlua_ref_is_function(LuaRef ref)
 /// @param name    if non-NULL, sent to callback as first arg
 ///                if NULL, only args are used
 /// @param retval  if true, convert return value to Object
-///                if false, discard return value
+///                if false, only check if return value is truthy
 /// @param err     Error details, if any (if NULL, errors are echoed)
-/// @return        Return value of function, if retval was set. Otherwise NIL.
+/// @return        Return value of function, if retval was set. Otherwise
+/// BOOLEAN_OBJ(true) or NIL.
 Object nlua_call_ref(LuaRef ref, const char *name, Array args, bool retval, Error *err)
 {
   lua_State *const lstate = global_lstate;
@@ -1438,7 +1516,7 @@ Object nlua_call_ref(LuaRef ref, const char *name, Array args, bool retval, Erro
     nlua_push_Object(lstate, args.items[i], false);
   }
 
-  if (nlua_pcall(lstate, nargs, retval ? 1 : 0)) {
+  if (nlua_pcall(lstate, nargs, 1)) {
     // if err is passed, the caller will deal with the error.
     if (err) {
       size_t len;
@@ -1458,7 +1536,10 @@ Object nlua_call_ref(LuaRef ref, const char *name, Array args, bool retval, Erro
     }
     return nlua_pop_Object(lstate, false, err);
   } else {
-    return NIL;
+    bool value = lua_toboolean(lstate, -1);
+    lua_pop(lstate, 1);
+
+    return value ? BOOLEAN_OBJ(true) : NIL;
   }
 }
 

src/nvim/map.c

@@ -168,6 +168,7 @@ MAP_IMPL(int, cstr_t, DEFAULT_INITIALIZER)
 MAP_IMPL(cstr_t, ptr_t, DEFAULT_INITIALIZER)
 MAP_IMPL(cstr_t, int, DEFAULT_INITIALIZER)
 MAP_IMPL(ptr_t, ptr_t, DEFAULT_INITIALIZER)
+MAP_IMPL(uint32_t, ptr_t, DEFAULT_INITIALIZER)
 MAP_IMPL(uint64_t, ptr_t, DEFAULT_INITIALIZER)
 MAP_IMPL(uint64_t, ssize_t, SSIZE_INITIALIZER)
 MAP_IMPL(uint64_t, uint64_t, DEFAULT_INITIALIZER)

src/nvim/map.h

@@ -39,6 +39,7 @@ MAP_DECLS(int, cstr_t)
 MAP_DECLS(cstr_t, ptr_t)
 MAP_DECLS(cstr_t, int)
 MAP_DECLS(ptr_t, ptr_t)
+MAP_DECLS(uint32_t, ptr_t)
 MAP_DECLS(uint64_t, ptr_t)
 MAP_DECLS(uint64_t, ssize_t)
 MAP_DECLS(uint64_t, uint64_t)

src/nvim/memory.c

@@ -23,6 +23,7 @@
 #include "nvim/message.h"
 #include "nvim/sign.h"
 #include "nvim/ui.h"
+#include "nvim/ui_compositor.h"
 #include "nvim/vim.h"
 
 #ifdef UNIT_TESTING
@@ -824,6 +825,7 @@ void free_all_mem(void)
 
   nlua_free_all_mem();
   ui_free_all_mem();
+  ui_comp_free_all_mem();
 
   // should be last, in case earlier free functions deallocates arenas
   arena_free_reuse_blks();

src/nvim/ui.c

@@ -198,13 +198,16 @@ void ui_refresh(void)
     ext_widgets[i] = true;
   }
 
+  UI *compositor = uis[0];
+
   bool inclusive = ui_override();
-  for (size_t i = 0; i < ui_count; i++) {
+  for (size_t i = 1; i < ui_count; i++) {
     UI *ui = uis[i];
     width = MIN(ui->width, width);
     height = MIN(ui->height, height);
     for (UIExtension j = 0; (int)j < kUIExtCount; j++) {
-      ext_widgets[j] &= (ui->ui_ext[j] || inclusive);
+      bool in_compositor = ui->composed && compositor->ui_ext[j];
+      ext_widgets[j] &= (ui->ui_ext[j] || in_compositor || inclusive);
     }
   }
 

src/nvim/ui.h

@@ -69,6 +69,11 @@ struct ui_t {
   void (*inspect)(UI *ui, Dictionary *info);
 };
 
+typedef struct ui_event_callback {
+  LuaRef cb;
+  bool ext_widgets[kUIGlobalCount];
+} UIEventCallback;
+
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "ui.h.generated.h"
 

src/nvim/ui_compositor.c

@@ -20,6 +20,7 @@
 #include "nvim/log.h"
 #include "nvim/lua/executor.h"
 #include "nvim/main.h"
+#include "nvim/map.h"
 #include "nvim/memory.h"
 #include "nvim/message.h"
 #include "nvim/os/os.h"
@@ -54,6 +55,8 @@ static bool msg_was_scrolled = false;
 static int msg_sep_row = -1;
 static schar_T msg_sep_char = { ' ', NUL };
 
+static PMap(uint32_t) ui_event_cbs = MAP_INIT;
+
 static int dbghl_normal, dbghl_clear, dbghl_composed, dbghl_recompose;
 
 void ui_comp_init(void)
@@ -69,14 +72,18 @@ void ui_comp_init(void)
   compositor->grid_cursor_goto = ui_comp_grid_cursor_goto;
   compositor->raw_line = ui_comp_raw_line;
   compositor->msg_set_pos = ui_comp_msg_set_pos;
+  compositor->event = ui_comp_event;
 
   // Be unopinionated: will be attached together with a "real" ui anyway
   compositor->width = INT_MAX;
   compositor->height = INT_MAX;
-  for (UIExtension i = 0; (int)i < kUIExtCount; i++) {
+  for (UIExtension i = kUIGlobalCount; (int)i < kUIExtCount; i++) {
     compositor->ui_ext[i] = true;
   }
 
+  // TODO(bfredl): one day. in the future.
+  compositor->ui_ext[kUIMultigrid] = false;
+
   // TODO(bfredl): this will be more complicated if we implement
   // hlstate per UI (i e reduce hl ids for non-hlstate UIs)
   compositor->ui_ext[kUIHlState] = false;
@@ -87,6 +94,15 @@ void ui_comp_init(void)
   ui_attach_impl(compositor, 0);
 }
 
+void ui_comp_free_all_mem(void)
+{
+  UIEventCallback *event_cb;
+  map_foreach_value(&ui_event_cbs, event_cb, {
+    xfree(event_cb);
+  })
+  pmap_destroy(uint32_t)(&ui_event_cbs);
+}
+
 void ui_comp_syn_init(void)
 {
   dbghl_normal = syn_check_group(S_LEN("RedrawDebugNormal"));
@@ -676,3 +692,72 @@ static void ui_comp_grid_resize(UI *ui, Integer grid, Integer width, Integer hei
     }
   }
 }
+
+static void ui_comp_event(UI *ui, char *name, Array args)
+{
+  Error err = ERROR_INIT;
+  UIEventCallback *event_cb;
+  bool handled = false;
+
+  map_foreach_value(&ui_event_cbs, event_cb, {
+    Object res = nlua_call_ref(event_cb->cb, name, args, false, &err);
+    if (res.type == kObjectTypeBoolean && res.data.boolean == true) {
+      handled = true;
+    }
+  })
+
+  if (!handled) {
+    ui_composed_call_event(name, args);
+  }
+}
+
+static void ui_comp_update_ext(void)
+{
+  memset(compositor->ui_ext, 0, ARRAY_SIZE(compositor->ui_ext));
+
+  for (size_t i = 0; i < kUIGlobalCount; i++) {
+    UIEventCallback *event_cb;
+
+    map_foreach_value(&ui_event_cbs, event_cb, {
+      if (event_cb->ext_widgets[i]) {
+        compositor->ui_ext[i] = true;
+        break;
+      }
+    })
+  }
+}
+
+void free_ui_event_callback(UIEventCallback *event_cb)
+{
+  api_free_luaref(event_cb->cb);
+  xfree(event_cb);
+}
+
+void ui_comp_add_cb(uint32_t ns_id, LuaRef cb, bool *ext_widgets)
+{
+  UIEventCallback *event_cb = xcalloc(1, sizeof(UIEventCallback));
+  event_cb->cb = cb;
+  memcpy(event_cb->ext_widgets, ext_widgets, ARRAY_SIZE(event_cb->ext_widgets));
+  if (event_cb->ext_widgets[kUIMessages]) {
+    event_cb->ext_widgets[kUICmdline] = true;
+  }
+
+  UIEventCallback **item = (UIEventCallback **)pmap_ref(uint32_t)(&ui_event_cbs, ns_id, true);
+  if (*item) {
+    free_ui_event_callback(*item);
+  }
+  *item = event_cb;
+
+  ui_comp_update_ext();
+  ui_schedule_refresh();
+}
+
+void ui_comp_remove_cb(uint32_t ns_id)
+{
+  if (pmap_has(uint32_t)(&ui_event_cbs, ns_id)) {
+    free_ui_event_callback(pmap_get(uint32_t)(&ui_event_cbs, ns_id));
+    pmap_del(uint32_t)(&ui_event_cbs, ns_id);
+  }
+  ui_comp_update_ext();
+  ui_schedule_refresh();
+}

test/functional/lua/ui_event_spec.lua

@@ -0,0 +1,107 @@
+local helpers = require('test.functional.helpers')(after_each)
+local Screen = require('test.functional.ui.screen')
+local eq = helpers.eq
+local exec_lua = helpers.exec_lua
+local clear = helpers.clear
+local feed = helpers.feed
+local funcs = helpers.funcs
+local inspect = require'vim.inspect'
+
+describe('vim.ui_attach', function()
+  local screen
+  before_each(function()
+    clear()
+    exec_lua [[
+      ns = vim.api.nvim_create_namespace 'testspace'
+      events = {}
+      function on_event(event, ...)
+        events[#events+1] = {event, ...}
+        return true
+      end
+
+      function get_events()
+        local ret_events = events
+        events = {}
+        return ret_events
+      end
+    ]]
+
+    screen = Screen.new(40,5)
+    screen:set_default_attr_ids({
+      [1] = {bold = true, foreground = Screen.colors.Blue1};
+      [2] = {bold = true};
+      [3] = {background = Screen.colors.Grey};
+      [4] = {background = Screen.colors.LightMagenta};
+    })
+    screen:attach()
+  end)
+
+  local function expect_events(expected)
+    local evs = exec_lua "return get_events(...)"
+    eq(expected, evs, inspect(evs))
+  end
+
+  it('can receive popupmenu events', function()
+    exec_lua [[ vim.ui_attach(ns, {ext_popupmenu=true}, on_event) ]]
+    feed('ifo')
+    screen:expect{grid=[[
+      fo^                                      |
+      {1:~                                       }|
+      {1:~                                       }|
+      {1:~                                       }|
+      {2:-- INSERT --}                            |
+    ]]}
+
+    funcs.complete(1, {'food', 'foobar', 'foo'})
+    screen:expect{grid=[[
+      food^                                    |
+      {1:~                                       }|
+      {1:~                                       }|
+      {1:~                                       }|
+      {2:-- INSERT --}                            |
+    ]]}
+    expect_events {
+      { "popupmenu_show", { { "food", "", "", "" }, { "foobar", "", "", "" }, { "foo", "", "", "" } }, 0, 0, 0, 1 };
+    }
+
+    feed '<c-n>'
+    screen:expect{grid=[[
+      foobar^                                  |
+      {1:~                                       }|
+      {1:~                                       }|
+      {1:~                                       }|
+      {2:-- INSERT --}                            |
+    ]]}
+    expect_events {
+       { "popupmenu_select", 1 };
+    }
+
+    feed '<c-y>'
+    screen:expect{grid=[[
+      foobar^                                  |
+      {1:~                                       }|
+      {1:~                                       }|
+      {1:~                                       }|
+      {2:-- INSERT --}                            |
+    ]], intermediate=true}
+    expect_events {
+       { "popupmenu_hide" };
+    }
+
+    -- ui_detach stops events, and reenables builtin pum
+    exec_lua [[ vim.ui_detach(ns) ]]
+
+    funcs.complete(1, {'food', 'foobar', 'foo'})
+    screen:expect{grid=[[
+      food^                                    |
+      {3:food           }{1:                         }|
+      {4:foobar         }{1:                         }|
+      {4:foo            }{1:                         }|
+      {2:-- INSERT --}                            |
+    ]]}
+    expect_events {
+    }
+
+
+  end)
+end)