feat(lua): add vim._with

It's a function to perform operations in their own sealed context,
similar to pythons `with`. This helps ease operations where you need to
perform an operation in a specific context, and then restore the
context.

Marked as private for now as it's not ready for public use. The current
plan is to start using this internally so we can discover and fix any
problems. Once this is ready to be exposed it will be renamed to
`vim.with`.

Usage:

```lua
local ret = vim._with({context = val}, function()
  return "hello"
end)
```

, where `context` is any combination of:

- `buf`
- `emsg_silent`
- `hide`
- `horizontal`
- `keepalt`
- `keepjumps`
- `keepmarks`
- `keeppatterns`
- `lockmarks`
- `noautocmd`
- `options`
- `sandbox`
- `silent`
- `unsilent`
- `win`

(except for `win` and `buf` which can't be used at the same time). This
list will most likely be expanded in the future.

Work on https://github.com/neovim/neovim/issues/19832.

Co-authored-by: Lewis Russell <lewis6991@gmail.com>

runtime/lua/vim/_comment.lua

@@ -194,14 +194,9 @@ local function toggle_lines(line_start, line_end, ref_position)
   -- - Debatable for highlighting in text area (like LSP semantic tokens).
   --   Mostly because it causes flicker as highlighting is preserved during
   --   comment toggling.
-  package.loaded['vim._comment']._lines = vim.tbl_map(f, lines)
-  local lua_cmd = string.format(
-    'vim.api.nvim_buf_set_lines(0, %d, %d, false, package.loaded["vim._comment"]._lines)',
-    line_start - 1,
-    line_end
-  )
-  vim.cmd.lua({ lua_cmd, mods = { lockmarks = true } })
-  package.loaded['vim._comment']._lines = nil
+  vim._with({ lockmarks = true }, function()
+    vim.api.nvim_buf_set_lines(0, line_start - 1, line_end, false, vim.tbl_map(f, lines))
+  end)
 end
 
 --- Operator which toggles user-supplied range of lines

runtime/lua/vim/shared.lua

@@ -1139,4 +1139,82 @@ function vim._defer_require(root, mod)
   })
 end
 
+--- @nodoc
+--- @class vim.context.mods
+--- @field buf? integer
+--- @field emsg_silent? boolean
+--- @field hide? boolean
+--- @field horizontal? boolean
+--- @field keepalt? boolean
+--- @field keepjumps? boolean
+--- @field keepmarks? boolean
+--- @field keeppatterns? boolean
+--- @field lockmarks? boolean
+--- @field noautocmd? boolean
+--- @field options? table<string, any>
+--- @field sandbox? boolean
+--- @field silent? boolean
+--- @field unsilent? boolean
+--- @field win? integer
+
+--- Executes function `f` with the given context specification.
+---
+--- @param context vim.context.mods
+function vim._with(context, f)
+  vim.validate('context', context, 'table')
+  vim.validate('f', f, 'function')
+
+  vim.validate('context.buf', context.buf, 'number', true)
+  vim.validate('context.emsg_silent', context.emsg_silent, 'boolean', true)
+  vim.validate('context.hide', context.hide, 'boolean', true)
+  vim.validate('context.horizontal', context.horizontal, 'boolean', true)
+  vim.validate('context.keepalt', context.keepalt, 'boolean', true)
+  vim.validate('context.keepjumps', context.keepjumps, 'boolean', true)
+  vim.validate('context.keepmarks', context.keepmarks, 'boolean', true)
+  vim.validate('context.keeppatterns', context.keeppatterns, 'boolean', true)
+  vim.validate('context.lockmarks', context.lockmarks, 'boolean', true)
+  vim.validate('context.noautocmd', context.noautocmd, 'boolean', true)
+  vim.validate('context.options', context.options, 'table', true)
+  vim.validate('context.sandbox', context.sandbox, 'boolean', true)
+  vim.validate('context.silent', context.silent, 'boolean', true)
+  vim.validate('context.unsilent', context.unsilent, 'boolean', true)
+  vim.validate('context.win', context.win, 'number', true)
+
+  -- Check buffer exists
+  if context.buf then
+    if not vim.api.nvim_buf_is_valid(context.buf) then
+      error('Invalid buffer id: ' .. context.buf)
+    end
+  end
+
+  -- Check window exists
+  if context.win then
+    if not vim.api.nvim_win_is_valid(context.win) then
+      error('Invalid window id: ' .. context.win)
+    end
+  end
+
+  -- Store original options
+  local previous_options ---@type table<string, any>
+  if context.options then
+    previous_options = {}
+    for k, v in pairs(context.options) do
+      previous_options[k] =
+        vim.api.nvim_get_option_value(k, { win = context.win, buf = context.buf })
+      vim.api.nvim_set_option_value(k, v, { win = context.win, buf = context.buf })
+    end
+  end
+
+  local retval = { vim._with_c(context, f) }
+
+  -- Restore original options
+  if previous_options then
+    for k, v in pairs(previous_options) do
+      vim.api.nvim_set_option_value(k, v, { win = context.win, buf = context.buf })
+    end
+  end
+
+  return unpack(retval)
+end
+
 return vim

src/nvim/ex_docmd.c

@@ -1729,12 +1729,6 @@ int execute_cmd(exarg_T *eap, CmdParseInfo *cmdinfo, bool preview)
   }
 
   const char *errormsg = NULL;
-#undef ERROR
-#define ERROR(msg) \
-  do { \
-    errormsg = msg; \
-    goto end; \
-  } while (0)
 
   cmdmod_T save_cmdmod = cmdmod;
   cmdmod = cmdinfo->cmdmod;
@@ -1745,16 +1739,19 @@ int execute_cmd(exarg_T *eap, CmdParseInfo *cmdinfo, bool preview)
   if (!MODIFIABLE(curbuf) && (eap->argt & EX_MODIFY)
       // allow :put in terminals
       && !(curbuf->terminal && eap->cmdidx == CMD_put)) {
-    ERROR(_(e_modifiable));
+    errormsg = _(e_modifiable);
+    goto end;
   }
   if (!IS_USER_CMDIDX(eap->cmdidx)) {
     if (cmdwin_type != 0 && !(eap->argt & EX_CMDWIN)) {
       // Command not allowed in the command line window
-      ERROR(_(e_cmdwin));
+      errormsg = _(e_cmdwin);
+      goto end;
     }
     if (text_locked() && !(eap->argt & EX_LOCK_OK)) {
       // Command not allowed when text is locked
-      ERROR(_(get_text_locked_msg()));
+      errormsg = _(get_text_locked_msg());
+      goto end;
     }
   }
   // Disallow editing another buffer when "curbuf->b_ro_locked" is set.
@@ -1802,7 +1799,6 @@ end:
 
   do_cmdline_end();
   return retv;
-#undef ERROR
 }
 
 static void profile_cmd(const exarg_T *eap, cstack_T *cstack, LineGetter fgetline, void *cookie)
@@ -2696,7 +2692,7 @@ int parse_command_modifiers(exarg_T *eap, const char **errormsg, cmdmod_T *cmod,
 
 /// Apply the command modifiers.  Saves current state in "cmdmod", call
 /// undo_cmdmod() later.
-static void apply_cmdmod(cmdmod_T *cmod)
+void apply_cmdmod(cmdmod_T *cmod)
 {
   if ((cmod->cmod_flags & CMOD_SANDBOX) && !cmod->cmod_did_sandbox) {
     sandbox++;

src/nvim/lua/stdlib.c

@@ -17,10 +17,13 @@
 #include "nvim/api/private/defs.h"
 #include "nvim/api/private/helpers.h"
 #include "nvim/ascii_defs.h"
+#include "nvim/autocmd.h"
 #include "nvim/buffer_defs.h"
 #include "nvim/eval/typval.h"
 #include "nvim/eval/typval_defs.h"
 #include "nvim/eval/vars.h"
+#include "nvim/eval/window.h"
+#include "nvim/ex_docmd.h"
 #include "nvim/ex_eval.h"
 #include "nvim/fold.h"
 #include "nvim/globals.h"
@@ -40,6 +43,7 @@
 #include "nvim/runtime.h"
 #include "nvim/strings.h"
 #include "nvim/types_defs.h"
+#include "nvim/window.h"
 
 #ifdef INCLUDE_GENERATED_DECLARATIONS
 # include "lua/stdlib.c.generated.h"
@@ -568,6 +572,99 @@ static int nlua_foldupdate(lua_State *lstate)
   return 0;
 }
 
+static int nlua_with(lua_State *L)
+{
+  int flags = 0;
+  buf_T *buf = NULL;
+  win_T *win = NULL;
+
+#define APPLY_FLAG(key, flag) \
+  if (strequal((key), k) && (v)) { \
+    flags |= (flag); \
+  }
+
+  luaL_argcheck(L, lua_istable(L, 1), 1, "table expected");
+  lua_pushnil(L);  // [dict, ..., nil]
+  while (lua_next(L, 1)) {
+    // [dict, ..., key, value]
+    if (lua_type(L, -2) == LUA_TSTRING) {
+      const char *k = lua_tostring(L, -2);
+      bool v = lua_toboolean(L, -1);
+      if (strequal("buf", k)) { \
+        buf = handle_get_buffer((int)luaL_checkinteger(L, -1));
+      } else if (strequal("win", k)) { \
+        win = handle_get_window((int)luaL_checkinteger(L, -1));
+      } else {
+        APPLY_FLAG("sandbox", CMOD_SANDBOX);
+        APPLY_FLAG("silent", CMOD_SILENT);
+        APPLY_FLAG("emsg_silent", CMOD_ERRSILENT);
+        APPLY_FLAG("unsilent", CMOD_UNSILENT);
+        APPLY_FLAG("noautocmd", CMOD_NOAUTOCMD);
+        APPLY_FLAG("hide", CMOD_HIDE);
+        APPLY_FLAG("keepalt", CMOD_KEEPALT);
+        APPLY_FLAG("keepmarks", CMOD_KEEPMARKS);
+        APPLY_FLAG("keepjumps", CMOD_KEEPJUMPS);
+        APPLY_FLAG("lockmarks", CMOD_LOCKMARKS);
+        APPLY_FLAG("keeppatterns", CMOD_KEEPPATTERNS);
+      }
+    }
+    // pop the value; lua_next will pop the key.
+    lua_pop(L, 1);  // [dict, ..., key]
+  }
+  int status = 0;
+  int rets = 0;
+
+  cmdmod_T save_cmdmod = cmdmod;
+  cmdmod.cmod_flags = flags;
+  apply_cmdmod(&cmdmod);
+
+  if (buf || win) {
+    try_start();
+  }
+
+  aco_save_T aco;
+  win_execute_T win_execute_args;
+  Error err = ERROR_INIT;
+
+  if (win) {
+    tabpage_T *tabpage = win_find_tabpage(win);
+    if (!win_execute_before(&win_execute_args, win, tabpage)) {
+      goto end;
+    }
+  } else if (buf) {
+    aucmd_prepbuf(&aco, buf);
+  }
+
+  int s = lua_gettop(L);
+  lua_pushvalue(L, 2);
+  status = lua_pcall(L, 0, LUA_MULTRET, 0);
+  rets = lua_gettop(L) - s;
+
+  if (win) {
+    win_execute_after(&win_execute_args);
+  } else if (buf) {
+    aucmd_restbuf(&aco);
+  }
+
+end:
+  if (buf || win) {
+    try_end(&err);
+  }
+
+  undo_cmdmod(&cmdmod);
+  cmdmod = save_cmdmod;
+
+  if (status) {
+    return lua_error(L);
+  } else if (ERROR_SET(&err)) {
+    nlua_push_errstr(L, "%s", err.msg);
+    api_clear_error(&err);
+    return lua_error(L);
+  }
+
+  return rets;
+}
+
 // Access to internal functions. For use in runtime/
 static void nlua_state_add_internal(lua_State *const lstate)
 {
@@ -582,6 +679,9 @@ static void nlua_state_add_internal(lua_State *const lstate)
   // _updatefolds
   lua_pushcfunction(lstate, &nlua_foldupdate);
   lua_setfield(lstate, -2, "_foldupdate");
+
+  lua_pushcfunction(lstate, &nlua_with);
+  lua_setfield(lstate, -2, "_with_c");
 }
 
 void nlua_state_add_stdlib(lua_State *const lstate, bool is_thread)

test/functional/lua/with_spec.lua

@@ -0,0 +1,292 @@
+local t = require('test.testutil')
+local n = require('test.functional.testnvim')()
+local Screen = require('test.functional.ui.screen')
+
+local fn = n.fn
+local api = n.api
+local command = n.command
+local eq = t.eq
+local exec_lua = n.exec_lua
+local matches = t.matches
+local pcall_err = t.pcall_err
+
+before_each(function()
+  n.clear()
+end)
+
+describe('vim._with {buf = }', function()
+  it('does not trigger autocmd', function()
+    exec_lua [[
+      local new = vim.api.nvim_create_buf(false, true)
+      vim.api.nvim_create_autocmd( { 'BufEnter', 'BufLeave', 'BufWinEnter', 'BufWinLeave' }, {
+        callback = function() _G.n = (_G.n or 0) + 1 end
+      })
+      vim._with({buf = new}, function()
+      end)
+      assert(_G.n == nil)
+    ]]
+  end)
+
+  it('trigger autocmd if changed within context', function()
+    exec_lua [[
+      local new = vim.api.nvim_create_buf(false, true)
+      vim.api.nvim_create_autocmd( { 'BufEnter', 'BufLeave', 'BufWinEnter', 'BufWinLeave' }, {
+        callback = function() _G.n = (_G.n or 0) + 1 end
+      })
+      vim._with({}, function()
+        vim.api.nvim_set_current_buf(new)
+        assert(_G.n ~= nil)
+      end)
+    ]]
+  end)
+
+  it('can access buf options', function()
+    local buf1 = api.nvim_get_current_buf()
+    local buf2 = exec_lua [[
+      buf2 = vim.api.nvim_create_buf(false, true)
+      return buf2
+    ]]
+
+    eq(false, api.nvim_get_option_value('autoindent', { buf = buf1 }))
+    eq(false, api.nvim_get_option_value('autoindent', { buf = buf2 }))
+
+    local val = exec_lua [[
+      return vim._with({buf = buf2}, function()
+        vim.cmd "set autoindent"
+        return vim.api.nvim_get_current_buf()
+    end)
+    ]]
+
+    eq(false, api.nvim_get_option_value('autoindent', { buf = buf1 }))
+    eq(true, api.nvim_get_option_value('autoindent', { buf = buf2 }))
+    eq(buf1, api.nvim_get_current_buf())
+    eq(buf2, val)
+  end)
+
+  it('does not cause ml_get errors with invalid visual selection', function()
+    exec_lua [[
+      local api = vim.api
+      local t = function(s) return api.nvim_replace_termcodes(s, true, true, true) end
+      api.nvim_buf_set_lines(0, 0, -1, true, {"a", "b", "c"})
+      api.nvim_feedkeys(t "G<C-V>", "txn", false)
+      vim._with({buf = api.nvim_create_buf(false, true)}, function() vim.cmd "redraw" end)
+    ]]
+  end)
+
+  it('can be nested crazily with hidden buffers', function()
+    eq(
+      true,
+      exec_lua([[
+      local function scratch_buf_call(fn)
+        local buf = vim.api.nvim_create_buf(false, true)
+        vim.api.nvim_set_option_value('cindent', true, {buf = buf})
+        return vim._with({buf = buf}, function()
+          return vim.api.nvim_get_current_buf() == buf
+            and vim.api.nvim_get_option_value('cindent', {buf = buf})
+            and fn()
+      end) and vim.api.nvim_buf_delete(buf, {}) == nil
+    end
+
+    return scratch_buf_call(function()
+      return scratch_buf_call(function()
+        return scratch_buf_call(function()
+          return scratch_buf_call(function()
+            return scratch_buf_call(function()
+              return scratch_buf_call(function()
+                return scratch_buf_call(function()
+                  return scratch_buf_call(function()
+                    return scratch_buf_call(function()
+                      return scratch_buf_call(function()
+                        return scratch_buf_call(function()
+                          return scratch_buf_call(function()
+                            return true
+                          end)
+                        end)
+                      end)
+                    end)
+                  end)
+                end)
+              end)
+            end)
+          end)
+        end)
+      end)
+    end)
+    ]])
+    )
+  end)
+
+  it('can return values by reference', function()
+    eq(
+      { 4, 7 },
+      exec_lua [[
+      local val = {4, 10}
+      local ref = vim._with({ buf = 0}, function() return val end)
+      ref[2] = 7
+      return val
+    ]]
+    )
+  end)
+end)
+
+describe('vim._with {win = }', function()
+  it('does not trigger autocmd', function()
+    exec_lua [[
+      local old = vim.api.nvim_get_current_win()
+      vim.cmd("new")
+      local new = vim.api.nvim_get_current_win()
+      vim.api.nvim_create_autocmd( { 'WinEnter', 'WinLeave' }, {
+        callback = function() _G.n = (_G.n or 0) + 1 end
+      })
+      vim._with({win = old}, function()
+      end)
+      assert(_G.n == nil)
+    ]]
+  end)
+
+  it('trigger autocmd if changed within context', function()
+    exec_lua [[
+      local old = vim.api.nvim_get_current_win()
+      vim.cmd("new")
+      local new = vim.api.nvim_get_current_win()
+      vim.api.nvim_create_autocmd( { 'WinEnter', 'WinLeave' }, {
+        callback = function() _G.n = (_G.n or 0) + 1 end
+      })
+      vim._with({}, function()
+        vim.api.nvim_set_current_win(old)
+        assert(_G.n ~= nil)
+      end)
+    ]]
+  end)
+
+  it('can access window options', function()
+    command('vsplit')
+    local win1 = api.nvim_get_current_win()
+    command('wincmd w')
+    local win2 = exec_lua [[
+      win2 = vim.api.nvim_get_current_win()
+      return win2
+    ]]
+    command('wincmd p')
+
+    eq('', api.nvim_get_option_value('winhighlight', { win = win1 }))
+    eq('', api.nvim_get_option_value('winhighlight', { win = win2 }))
+
+    local val = exec_lua [[
+      return vim._with({win = win2}, function()
+        vim.cmd "setlocal winhighlight=Normal:Normal"
+        return vim.api.nvim_get_current_win()
+      end)
+    ]]
+
+    eq('', api.nvim_get_option_value('winhighlight', { win = win1 }))
+    eq('Normal:Normal', api.nvim_get_option_value('winhighlight', { win = win2 }))
+    eq(win1, api.nvim_get_current_win())
+    eq(win2, val)
+  end)
+
+  it('does not cause ml_get errors with invalid visual selection', function()
+    -- Add lines to the current buffer and make another window looking into an empty buffer.
+    exec_lua [[
+      _G.api = vim.api
+      _G.t = function(s) return api.nvim_replace_termcodes(s, true, true, true) end
+      _G.win_lines = api.nvim_get_current_win()
+      vim.cmd "new"
+      _G.win_empty = api.nvim_get_current_win()
+      api.nvim_set_current_win(win_lines)
+      api.nvim_buf_set_lines(0, 0, -1, true, {"a", "b", "c"})
+    ]]
+
+    -- Start Visual in current window, redraw in other window with fewer lines.
+    exec_lua [[
+      api.nvim_feedkeys(t "G<C-V>", "txn", false)
+      vim._with({win = win_empty}, function() vim.cmd "redraw" end)
+    ]]
+
+    -- Start Visual in current window, extend it in other window with more lines.
+    exec_lua [[
+      api.nvim_feedkeys(t "<Esc>gg", "txn", false)
+      api.nvim_set_current_win(win_empty)
+      api.nvim_feedkeys(t "gg<C-V>", "txn", false)
+      vim._with({win = win_lines}, function() api.nvim_feedkeys(t "G<C-V>", "txn", false) end)
+      vim.cmd "redraw"
+    ]]
+  end)
+
+  it('updates ruler if cursor moved', function()
+    local screen = Screen.new(30, 5)
+    screen:set_default_attr_ids {
+      [1] = { reverse = true },
+      [2] = { bold = true, reverse = true },
+    }
+    screen:attach()
+    exec_lua [[
+      _G.api = vim.api
+      vim.opt.ruler = true
+      local lines = {}
+      for i = 0, 499 do lines[#lines + 1] = tostring(i) end
+      api.nvim_buf_set_lines(0, 0, -1, true, lines)
+      api.nvim_win_set_cursor(0, {20, 0})
+      vim.cmd "split"
+      _G.win = api.nvim_get_current_win()
+      vim.cmd "wincmd w | redraw"
+    ]]
+    screen:expect [[
+      19                            |
+      {1:[No Name] [+]  20,1         3%}|
+      ^19                            |
+      {2:[No Name] [+]  20,1         3%}|
+                                    |
+    ]]
+    exec_lua [[
+      vim._with({win = win}, function() api.nvim_win_set_cursor(0, {100, 0}) end)
+      vim.cmd "redraw"
+    ]]
+    screen:expect [[
+      99                            |
+      {1:[No Name] [+]  100,1       19%}|
+      ^19                            |
+      {2:[No Name] [+]  20,1         3%}|
+                                    |
+    ]]
+  end)
+
+  it('can return values by reference', function()
+    eq(
+      { 7, 10 },
+      exec_lua [[
+      local val = {4, 10}
+      local ref = vim._with({win = 0}, function() return val end)
+      ref[1] = 7
+      return val
+    ]]
+    )
+  end)
+
+  it('layout in current tabpage does not affect windows in others', function()
+    command('tab split')
+    local t2_move_win = api.nvim_get_current_win()
+    command('vsplit')
+    local t2_other_win = api.nvim_get_current_win()
+    command('tabprevious')
+    matches('E36: Not enough room$', pcall_err(command, 'execute "split|"->repeat(&lines)'))
+    command('vsplit')
+
+    exec_lua('vim._with({win = ...}, function() vim.cmd.wincmd "J" end)', t2_move_win)
+    eq({ 'col', { { 'leaf', t2_other_win }, { 'leaf', t2_move_win } } }, fn.winlayout(2))
+  end)
+end)
+
+describe('vim._with {lockmarks = true}', function()
+  it('is reset', function()
+    local mark = exec_lua [[
+      vim.api.nvim_buf_set_lines(0, 0, 0, false, {"marky", "snarky", "malarkey"})
+      vim.api.nvim_buf_set_mark(0,"m",1,0, {})
+      vim._with({lockmarks = true}, function()
+        vim.api.nvim_buf_set_lines(0, 0, 2, false, {"mass", "mess", "moss"})
+      end)
+      return vim.api.nvim_buf_get_mark(0,"m")
+    ]]
+    t.eq(mark, { 1, 0 })
+  end)
+end)