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>

5 files changed, 480 insertions, 19 deletions

diff --git a/runtime/lua/vim/_comment.lua b/runtime/lua/vim/_comment.lua
index 044cd69716..efe289b3e1 100644
--- a/runtime/lua/vim/_comment.lua
+++ b/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
diff --git a/runtime/lua/vim/shared.lua b/runtime/lua/vim/shared.lua
index 9d54a80144..7fd29d5f7b 100644
--- a/runtime/lua/vim/shared.lua
+++ b/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
diff --git a/src/nvim/ex_docmd.c b/src/nvim/ex_docmd.c
index 2f54aa511b..1e2c515195 100644
--- a/src/nvim/ex_docmd.c
+++ b/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++;
diff --git a/src/nvim/lua/stdlib.c b/src/nvim/lua/stdlib.c
index 22ee0a1c98..ee0eabbebb 100644
--- a/src/nvim/lua/stdlib.c
+++ b/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)
diff --git a/test/functional/lua/with_spec.lua b/test/functional/lua/with_spec.lua
new file mode 100644
index 0000000000..36dee9630a
--- /dev/null
+++ b/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)