1 files changed, 262 insertions, 35 deletions

diff --git a/runtime/lua/vim/shared.lua b/runtime/lua/vim/shared.lua
index e9e4326057..4d06cdd77d 100644
--- a/runtime/lua/vim/shared.lua
+++ b/runtime/lua/vim/shared.lua
@@ -214,7 +214,7 @@ end
 ---@param t table<T, any> (table) Table
 ---@return T[] : List of keys
 function vim.tbl_keys(t)
-  vim.validate({ t = { t, 't' } })
+  vim.validate('t', t, 'table')
   --- @cast t table<any,any>
 
   local keys = {}
@@ -231,7 +231,7 @@ end
 ---@param t table<any, T> (table) Table
 ---@return T[] : List of values
 function vim.tbl_values(t)
-  vim.validate({ t = { t, 't' } })
+  vim.validate('t', t, 'table')
 
   local values = {}
   for _, v in
@@ -332,7 +332,7 @@ end
 ---@param value any Value to compare
 ---@return boolean `true` if `t` contains `value`
 function vim.list_contains(t, value)
-  vim.validate({ t = { t, 't' } })
+  vim.validate('t', t, 'table')
   --- @cast t table<any,any>
 
   for _, v in ipairs(t) do
@@ -350,41 +350,32 @@ end
 ---@param t table Table to check
 ---@return boolean `true` if `t` is empty
 function vim.tbl_isempty(t)
-  vim.validate({ t = { t, 't' } })
+  vim.validate('t', t, 'table')
   return next(t) == nil
 end
 
---- We only merge empty tables or tables that are not an array (indexed by integers)
+--- We only merge empty tables or tables that are not list-like (indexed by consecutive integers
+--- starting from 1)
 local function can_merge(v)
-  return type(v) == 'table' and (vim.tbl_isempty(v) or not vim.isarray(v))
+  return type(v) == 'table' and (vim.tbl_isempty(v) or not vim.islist(v))
 end
 
-local function tbl_extend(behavior, deep_extend, ...)
-  if behavior ~= 'error' and behavior ~= 'keep' and behavior ~= 'force' then
-    error('invalid "behavior": ' .. tostring(behavior))
-  end
-
-  if select('#', ...) < 2 then
-    error(
-      'wrong number of arguments (given '
-        .. tostring(1 + select('#', ...))
-        .. ', expected at least 3)'
-    )
-  end
-
+--- Recursive worker for tbl_extend
+--- @param behavior 'error'|'keep'|'force'
+--- @param deep_extend boolean
+--- @param ... table<any,any>
+local function tbl_extend_rec(behavior, deep_extend, ...)
   local ret = {} --- @type table<any,any>
   if vim._empty_dict_mt ~= nil and getmetatable(select(1, ...)) == vim._empty_dict_mt then
     ret = vim.empty_dict()
   end
 
   for i = 1, select('#', ...) do
-    local tbl = select(i, ...)
-    vim.validate({ ['after the second argument'] = { tbl, 't' } })
-    --- @cast tbl table<any,any>
+    local tbl = select(i, ...) --[[@as table<any,any>]]
     if tbl then
       for k, v in pairs(tbl) do
         if deep_extend and can_merge(v) and can_merge(ret[k]) then
-          ret[k] = tbl_extend(behavior, true, ret[k], v)
+          ret[k] = tbl_extend_rec(behavior, true, ret[k], v)
         elseif behavior ~= 'force' and ret[k] ~= nil then
           if behavior == 'error' then
             error('key found in more than one map: ' .. k)
@@ -395,9 +386,31 @@ local function tbl_extend(behavior, deep_extend, ...)
       end
     end
   end
+
   return ret
 end
 
+--- @param behavior 'error'|'keep'|'force'
+--- @param deep_extend boolean
+--- @param ... table<any,any>
+local function tbl_extend(behavior, deep_extend, ...)
+  if behavior ~= 'error' and behavior ~= 'keep' and behavior ~= 'force' then
+    error('invalid "behavior": ' .. tostring(behavior))
+  end
+
+  local nargs = select('#', ...)
+
+  if nargs < 2 then
+    error(('wrong number of arguments (given %d, expected at least 3)'):format(1 + nargs))
+  end
+
+  for i = 1, nargs do
+    vim.validate('after the second argument', select(i, ...), 'table')
+  end
+
+  return tbl_extend_rec(behavior, deep_extend, ...)
+end
+
 --- Merges two or more tables.
 ---
 ---@see |extend()|
@@ -414,6 +427,11 @@ end
 
 --- Merges recursively two or more tables.
 ---
+--- Only values that are empty tables or tables that are not |lua-list|s (indexed by consecutive
+--- integers starting from 1) are merged recursively. This is useful for merging nested tables
+--- like default and user configurations where lists should be treated as literals (i.e., are
+--- overwritten instead of merged).
+---
 ---@see |vim.tbl_extend()|
 ---
 ---@generic T1: table
@@ -580,7 +598,7 @@ end
 ---@return fun(table: table<K, V>, index?: K):K, V # |for-in| iterator over sorted keys and their values
 ---@return T
 function vim.spairs(t)
-  assert(type(t) == 'table', ('expected table, got %s'):format(type(t)))
+  vim.validate('t', t, 'table')
   --- @cast t table<any,any>
 
   -- collect the keys
@@ -691,7 +709,7 @@ end
 ---@param t table Table
 ---@return integer : Number of non-nil values in table
 function vim.tbl_count(t)
-  vim.validate({ t = { t, 't' } })
+  vim.validate('t', t, 'table')
   --- @cast t table<any,any>
 
   local count = 0
@@ -723,7 +741,7 @@ end
 ---@param s string String to trim
 ---@return string String with whitespace removed from its beginning and end
 function vim.trim(s)
-  vim.validate({ s = { s, 's' } })
+  vim.validate('s', s, 'string')
   return s:match('^%s*(.*%S)') or ''
 end
 
@@ -733,7 +751,7 @@ end
 ---@param s string String to escape
 ---@return string %-escaped pattern string
 function vim.pesc(s)
-  vim.validate({ s = { s, 's' } })
+  vim.validate('s', s, 'string')
   return (s:gsub('[%(%)%.%%%+%-%*%?%[%]%^%$]', '%%%1'))
 end
 
@@ -743,7 +761,8 @@ end
 ---@param prefix string Prefix to match
 ---@return boolean `true` if `prefix` is a prefix of `s`
 function vim.startswith(s, prefix)
-  vim.validate({ s = { s, 's' }, prefix = { prefix, 's' } })
+  vim.validate('s', s, 'string')
+  vim.validate('prefix', prefix, 'string')
   return s:sub(1, #prefix) == prefix
 end
 
@@ -753,7 +772,8 @@ end
 ---@param suffix string Suffix to match
 ---@return boolean `true` if `suffix` is a suffix of `s`
 function vim.endswith(s, suffix)
-  vim.validate({ s = { s, 's' }, suffix = { suffix, 's' } })
+  vim.validate('s', s, 'string')
+  vim.validate('suffix', suffix, 'string')
   return #suffix == 0 or s:sub(-#suffix) == suffix
 end
 
@@ -787,7 +807,7 @@ do
   }
 
   --- @nodoc
-  --- @class vim.validate.Spec {[1]: any, [2]: string|string[], [3]: boolean }
+  --- @class vim.validate.Spec [any, string|string[], boolean]
   --- @field [1] any Argument value
   --- @field [2] string|string[]|fun(v:any):boolean, string? Type name, or callable
   --- @field [3]? boolean
@@ -877,8 +897,30 @@ do
     return true
   end
 
-  --- Validates a parameter specification (types and values). Specs are evaluated in alphanumeric
-  --- order, until the first failure.
+  --- Validate function arguments.
+  ---
+  --- This function has two valid forms:
+  ---
+  --- 1. vim.validate(name: str, value: any, type: string, optional?: bool)
+  --- 2. vim.validate(spec: table)
+  ---
+  --- Form 1 validates that argument {name} with value {value} has the type
+  --- {type}. {type} must be a value returned by |lua-type()|. If {optional} is
+  --- true, then {value} may be null. This form is significantly faster and
+  --- should be preferred for simple cases.
+  ---
+  --- Example:
+  ---
+  --- ```lua
+  --- function vim.startswith(s, prefix)
+  ---   vim.validate('s', s, 'string')
+  ---   vim.validate('prefix', prefix, 'string')
+  ---   ...
+  --- end
+  --- ```
+  ---
+  --- Form 2 validates a parameter specification (types and values). Specs are
+  --- evaluated in alphanumeric order, until the first failure.
   ---
   --- Usage example:
   ---
@@ -930,8 +972,32 @@ do
   ---               only if the argument is valid. Can optionally return an additional
   ---               informative error message as the second returned value.
   ---             - msg: (optional) error string if validation fails
-  function vim.validate(opt)
-    local ok, err_msg = is_valid(opt)
+  --- @overload fun(name: string, val: any, expected: string, optional?: boolean)
+  function vim.validate(opt, ...)
+    local ok = false
+    local err_msg ---@type string?
+    local narg = select('#', ...)
+    if narg == 0 then
+      ok, err_msg = is_valid(opt)
+    elseif narg >= 2 then
+      -- Overloaded signature for fast/simple cases
+      local name = opt --[[@as string]]
+      local v, expected, optional = ... ---@type string, string, boolean?
+      local actual = type(v)
+
+      ok = (actual == expected) or (v == nil and optional == true)
+      if not ok then
+        err_msg = ('%s: expected %s, got %s%s'):format(
+          name,
+          expected,
+          actual,
+          v and (' (%s)'):format(v) or ''
+        )
+      end
+    else
+      error('invalid arguments')
+    end
+
     if not ok then
       error(err_msg, 2)
     end
@@ -949,7 +1015,7 @@ function vim.is_callable(f)
   if m == nil then
     return false
   end
-  return type(m.__call) == 'function'
+  return type(rawget(m, '__call')) == 'function'
 end
 
 --- Creates a table whose missing keys are provided by {createfn} (like Python's "defaultdict").
@@ -1091,4 +1157,165 @@ function vim._defer_require(root, mod)
   })
 end
 
+--- @nodoc
+--- @class vim.context.mods
+--- @field bo? table<string, any>
+--- @field buf? integer
+--- @field emsg_silent? boolean
+--- @field env? table<string, any>
+--- @field go? table<string, any>
+--- @field hide? boolean
+--- @field keepalt? boolean
+--- @field keepjumps? boolean
+--- @field keepmarks? boolean
+--- @field keeppatterns? boolean
+--- @field lockmarks? boolean
+--- @field noautocmd? boolean
+--- @field o? table<string, any>
+--- @field sandbox? boolean
+--- @field silent? boolean
+--- @field unsilent? boolean
+--- @field win? integer
+--- @field wo? table<string, any>
+
+--- @nodoc
+--- @class vim.context.state
+--- @field bo? table<string, any>
+--- @field env? table<string, any>
+--- @field go? table<string, any>
+--- @field wo? table<string, any>
+
+local scope_map = { buf = 'bo', global = 'go', win = 'wo' }
+local scope_order = { 'o', 'wo', 'bo', 'go', 'env' }
+local state_restore_order = { 'bo', 'wo', 'go', 'env' }
+
+--- Gets data about current state, enough to properly restore specified options/env/etc.
+--- @param context vim.context.mods
+--- @return vim.context.state
+local get_context_state = function(context)
+  local res = { bo = {}, env = {}, go = {}, wo = {} }
+
+  -- Use specific order from possibly most to least intrusive
+  for _, scope in ipairs(scope_order) do
+    for name, _ in pairs(context[scope] or {}) do
+      local sc = scope == 'o' and scope_map[vim.api.nvim_get_option_info2(name, {}).scope] or scope
+
+      -- Do not override already set state and fall back to `vim.NIL` for
+      -- state `nil` values (which still needs restoring later)
+      res[sc][name] = res[sc][name] or vim[sc][name] or vim.NIL
+
+      -- Always track global option value to properly restore later.
+      -- This matters for at least `o` and `wo` (which might set either/both
+      -- local and global option values).
+      if sc ~= 'env' then
+        res.go[name] = res.go[name] or vim.go[name]
+      end
+    end
+  end
+
+  return res
+end
+
+--- Executes function `f` with the given context specification.
+---
+--- Notes:
+--- - Context `{ buf = buf }` has no guarantees about current window when
+---   inside context.
+--- - Context `{ buf = buf, win = win }` is yet not allowed, but this seems
+---   to be an implementation detail.
+--- - There should be no way to revert currently set `context.sandbox = true`
+---   (like with nested `vim._with()` calls). Otherwise it kind of breaks the
+---   whole purpose of sandbox execution.
+--- - Saving and restoring option contexts (`bo`, `go`, `o`, `wo`) trigger
+---   `OptionSet` events. This is an implementation issue because not doing it
+---   seems to mean using either 'eventignore' option or extra nesting with
+---   `{ noautocmd = true }` (which itself is a wrapper for 'eventignore').
+---   As `{ go = { eventignore = '...' } }` is a valid context which should be
+---   properly set and restored, this is not a good approach.
+---   Not triggering `OptionSet` seems to be a good idea, though. So probably
+---   only moving context save and restore to lower level might resolve this.
+---
+--- @param context vim.context.mods
+--- @param f function
+--- @return any
+function vim._with(context, f)
+  vim.validate('context', context, 'table')
+  vim.validate('f', f, 'function')
+
+  vim.validate('context.bo', context.bo, 'table', true)
+  vim.validate('context.buf', context.buf, 'number', true)
+  vim.validate('context.emsg_silent', context.emsg_silent, 'boolean', true)
+  vim.validate('context.env', context.env, 'table', true)
+  vim.validate('context.go', context.go, 'table', true)
+  vim.validate('context.hide', context.hide, '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.o', context.o, '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)
+  vim.validate('context.wo', context.wo, 'table', 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
+    -- TODO: Maybe allow it?
+    if context.buf and vim.api.nvim_win_get_buf(context.win) ~= context.buf then
+      error('Can not set both `buf` and `win` context.')
+    end
+  end
+
+  -- Decorate so that save-set-restore options is done in correct window-buffer
+  local callback = function()
+    -- Cache current values to be changed by context
+    -- Abort early in case of bad context value
+    local ok, state = pcall(get_context_state, context)
+    if not ok then
+      error(state, 0)
+    end
+
+    -- Apply some parts of the context in specific order
+    -- NOTE: triggers `OptionSet` event
+    for _, scope in ipairs(scope_order) do
+      for name, context_value in pairs(context[scope] or {}) do
+        vim[scope][name] = context_value
+      end
+    end
+
+    -- Execute
+    local res = { pcall(f) }
+
+    -- Restore relevant cached values in specific order, global scope last
+    -- NOTE: triggers `OptionSet` event
+    for _, scope in ipairs(state_restore_order) do
+      for name, cached_value in pairs(state[scope]) do
+        vim[scope][name] = cached_value
+      end
+    end
+
+    -- Return
+    if not res[1] then
+      error(res[2], 0)
+    end
+    table.remove(res, 1)
+    return unpack(res, 1, table.maxn(res))
+  end
+
+  return vim._with_c(context, callback)
+end
+
 return vim