feat(lua): function `behavior` for tbl_extend, tbl_deep_extend #33819

1 files changed, 16 insertions, 5 deletions

diff --git a/runtime/lua/vim/shared.lua b/runtime/lua/vim/shared.lua
index 16414e2828..18195d1126 100644
--- a/runtime/lua/vim/shared.lua
+++ b/runtime/lua/vim/shared.lua
@@ -366,7 +366,7 @@ local function can_merge(v)
 end
 
 --- Recursive worker for tbl_extend
---- @param behavior 'error'|'keep'|'force'
+--- @param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any): any
 --- @param deep_extend boolean
 --- @param ... table<any,any>
 local function tbl_extend_rec(behavior, deep_extend, ...)
@@ -381,6 +381,8 @@ local function tbl_extend_rec(behavior, deep_extend, ...)
       for k, v in pairs(tbl) do
         if deep_extend and can_merge(v) and can_merge(ret[k]) then
           ret[k] = tbl_extend_rec(behavior, true, ret[k], v)
+        elseif type(behavior) == 'function' then
+          ret[k] = behavior(k, v, ret[k])
         elseif behavior ~= 'force' and ret[k] ~= nil then
           if behavior == 'error' then
             error('key found in more than one map: ' .. k)
@@ -395,11 +397,16 @@ local function tbl_extend_rec(behavior, deep_extend, ...)
   return ret
 end
 
---- @param behavior 'error'|'keep'|'force'
+--- @param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any): any
 --- @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
+  if
+    behavior ~= 'error'
+    and behavior ~= 'keep'
+    and behavior ~= 'force'
+    and type(behavior) ~= 'function'
+  then
     error('invalid "behavior": ' .. tostring(behavior))
   end
 
@@ -420,10 +427,12 @@ end
 ---
 ---@see |extend()|
 ---
----@param behavior 'error'|'keep'|'force' Decides what to do if a key is found in more than one map:
+---@param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any?): any Decides what to do if a key is found in more than one map:
 ---      - "error": raise an error
 ---      - "keep":  use value from the leftmost map
 ---      - "force": use value from the rightmost map
+---      - If a function, it receives the current key, value, and the previous value in the currently merged table (if present) and should
+---        return the value for the given key in the merged table.
 ---@param ... table Two or more tables
 ---@return table : Merged table
 function vim.tbl_extend(behavior, ...)
@@ -441,10 +450,12 @@ end
 ---
 ---@generic T1: table
 ---@generic T2: table
----@param behavior 'error'|'keep'|'force' Decides what to do if a key is found in more than one map:
+---@param behavior 'error'|'keep'|'force'|fun(key:any, v:any, prev_value:any?): any Decides what to do if a key is found in more than one map:
 ---      - "error": raise an error
 ---      - "keep":  use value from the leftmost map
 ---      - "force": use value from the rightmost map
+---      - If a function, it receives the current key, value, and the previous value in the currently merged table (if present) and should
+---        return the value for the given key in the merged table.
 ---@param ... T2 Two or more tables
 ---@return T1|T2 (table) Merged table
 function vim.tbl_deep_extend(behavior, ...)