perf(lua): memoize key_fn results #39568

Problem:
When using `vim.list.unique` or `vim.list.bisect`, if the `key` function is
complex, it can degrade performance, because it is invoked on every comparison

Solution:
The `key` interface convention is designed specifically to address this issue;
performance can be improved by memoizing its results.

Also added the shorthand use of the field name string as the key.

runtime/doc/lua.txt

@@ -1816,8 +1816,10 @@ vim.list.bisect({t}, {val}, {opts})                        *vim.list.bisect()*
                     keeps the list sorted..
                 • {hi}? (`integer`, default: `#t + 1`) End index of the list,
                   exclusive.
-                • {key}? (`fun(val: any): any`) Optional, compare the return
-                  value instead of the {val} itself if provided.
+                • {key}? (`string|fun(val: any): any`) Optional, compare the
+                  return value instead of the {val} itself if provided. If a
+                  string, index each value by this field name. Key results are
+                  memoized per call.
                 • {lo}? (`integer`, default: `1`) Start index of the list.
 
     Return: ~
@@ -1831,10 +1833,11 @@ vim.list.unique({t}, {key})                                *vim.list.unique()*
     performed in-place and the input table is modified.
 
     Accepts an optional `key` argument, which if provided is called for each
-    value in the list to compute a hash key for uniqueness comparison. This is
-    useful for deduplicating table values or complex objects. If `key` returns
-    `nil` for a value, that value will be considered unique, even if multiple
-    values return `nil`.
+    value in the list to compute a hash key for uniqueness comparison. If
+    `key` is a string, it is used as the field name to index each value. Key
+    results are memoized per call. This is useful for deduplicating table
+    values or complex objects. If `key` returns `nil` for a value, that value
+    will be considered unique, even if multiple values return `nil`.
 
     Example: >lua
 
@@ -1843,7 +1846,7 @@ vim.list.unique({t}, {key})                                *vim.list.unique()*
         -- t is now {1, 2, 3}
 
         local t = { {id=1}, {id=2}, {id=1} }
-        vim.list.unique(t, function(x) return x.id end)
+        vim.list.unique(t, 'id')
         -- t is now { {id=1}, {id=2} }
 <
 
@@ -1852,8 +1855,8 @@ vim.list.unique({t}, {key})                                *vim.list.unique()*
 
     Parameters: ~
       • {t}    (`any[]`)
-      • {key}  (`fun(x: T): any?`) Optional hash function to determine
-               uniqueness of values
+      • {key}  (`string|fun(x: T): any?`) Optional field name or hash function
+               to determine uniqueness of values
 
     Return: ~
         (`any[]`) The deduplicated list

runtime/doc/news.txt

@@ -172,6 +172,8 @@ LUA
   the same range instances now compare equal.
 • Documentation for |vim.fn| now points to preferred Lua alternatives when
   available.
+• |vim.list.unique()| and |vim.list.bisect()| now support passing a string
+  as a shorthand of a `key`
 
 OPTIONS
 
@@ -194,6 +196,8 @@ PERFORMANCE
   thus reducing GC and memory reallocation during each data reset.
 • When parsing the received Content-Length messages,
   the RPC client will no longer allocate extra strings.
+• |vim.list.unique()| and |vim.list.bisect()| now memoize key function results,
+  which can speed up calls with expensive key functions.
 
 PLUGINS
 

runtime/lua/vim/_core/shared.lua

@@ -376,13 +376,45 @@ end
 
 vim.list = {}
 
----TODO(ofseed): memoize, string value support, type alias.
+--- Returns a `key` function with per-call memoization.
 ---@generic T
----@param v T
----@param key? fun(v: T): any
----@return any
-local function key_fn(v, key)
-  return key and key(v) or v
+---@param key? string|fun(val: T): any
+---@return fun(v: T): any
+local function make_key_fn(key)
+  vim.validate('key', key, { 'string', 'function' }, true)
+
+  if not key then
+    return function(v)
+      return v
+    end
+  end
+
+  if type(key) == 'string' then
+    local field = key
+    ---@param v any
+    key = function(v)
+      return v and v[field] or nil
+    end
+  end
+
+  -- Keep memoized keys local to one list operation to avoid stale results.
+  local cache = {} --- @type table<any,any>
+
+  return function(v)
+    if v == nil then
+      return key(v)
+    end
+
+    local cached = cache[v]
+    if cached ~= nil then
+      -- Use `vim.NIL` to remember that `key(v)` returned `nil`.
+      return cached == vim.NIL and nil or cached
+    end
+
+    local result = key(v)
+    cache[v] = result == nil and vim.NIL or result
+    return result
+  end
 end
 
 --- Removes duplicate values from a |lua-list| in-place.
@@ -392,6 +424,8 @@ end
 ---
 --- Accepts an optional `key` argument, which if provided is called for each
 --- value in the list to compute a hash key for uniqueness comparison.
+--- If `key` is a string, it is used as the field name to index each value.
+--- Key results are memoized per call.
 --- This is useful for deduplicating table values or complex objects.
 --- If `key` returns `nil` for a value, that value will be considered unique,
 --- even if multiple values return `nil`.
@@ -404,18 +438,19 @@ end
 --- -- t is now {1, 2, 3}
 ---
 --- local t = { {id=1}, {id=2}, {id=1} }
---- vim.list.unique(t, function(x) return x.id end)
+--- vim.list.unique(t, 'id')
 --- -- t is now { {id=1}, {id=2} }
 --- ```
 ---
 --- @since 14
 --- @generic T
 --- @param t T[]
---- @param key? fun(x: T): any Optional hash function to determine uniqueness of values
+--- @param key? string|fun(x: T): any Optional field name or hash function to determine uniqueness of values
 --- @return T[] : The deduplicated list
 --- @see |Iter:unique()|
 function vim.list.unique(t, key)
   vim.validate('t', t, 'table')
+  local key_fn = make_key_fn(key)
   local seen = {} --- @type table<any,boolean>
 
   local finish = #t
@@ -423,7 +458,7 @@ function vim.list.unique(t, key)
   local j = 1
   for i = 1, finish do
     local v = t[i]
-    local vh = key_fn(v, key)
+    local vh = key_fn(v)
     if not seen[vh] then
       t[j] = v
       if vh ~= nil then
@@ -452,7 +487,9 @@ end
 ---@field hi? integer
 ---
 --- Optional, compare the return value instead of the {val} itself if provided.
----@field key? fun(val: any): any
+--- If a string, index each value by this field name.
+--- Key results are memoized per call.
+---@field key? string|fun(val: any): any
 ---
 --- Specifies the search variant.
 ---   - "lower": returns the first position
@@ -465,18 +502,18 @@ end
 ---@generic T
 ---@param t T[]
 ---@param val T
----@param key? fun(val: any): any
 ---@param lo integer
 ---@param hi integer
+---@param key_fn fun(val: any): any
 ---@return integer i in range such that `t[j]` < {val} for all j < i,
 ---                and `t[j]` >= {val} for all j >= i,
 ---                or return {hi} if no such index is found.
-local function lower_bound(t, val, lo, hi, key)
+local function lower_bound(t, val, lo, hi, key_fn)
   local bit = require('bit') -- Load bitop on demand
-  local val_key = key_fn(val, key)
+  local val_key = key_fn(val)
   while lo < hi do
     local mid = bit.rshift(lo + hi, 1) -- Equivalent to floor((lo + hi) / 2)
-    if key_fn(t[mid], key) < val_key then
+    if key_fn(t[mid]) < val_key then
       lo = mid + 1
     else
       hi = mid
@@ -488,18 +525,18 @@ end
 ---@generic T
 ---@param t T[]
 ---@param val T
----@param key? fun(val: any): any
 ---@param lo integer
 ---@param hi integer
+---@param key_fn fun(val: any): any
 ---@return integer i in range such that `t[j]` <= {val} for all j < i,
 ---                and `t[j]` > {val} for all j >= i,
 ---                or return {hi} if no such index is found.
-local function upper_bound(t, val, lo, hi, key)
+local function upper_bound(t, val, lo, hi, key_fn)
   local bit = require('bit') -- Load bitop on demand
-  local val_key = key_fn(val, key)
+  local val_key = key_fn(val)
   while lo < hi do
     local mid = bit.rshift(lo + hi, 1) -- Equivalent to floor((lo + hi) / 2)
-    if val_key < key_fn(t[mid], key) then
+    if val_key < key_fn(t[mid]) then
       hi = mid
     else
       lo = mid + 1
@@ -553,12 +590,12 @@ function vim.list.bisect(t, val, opts)
   opts = opts or {}
   local lo = opts.lo or 1
   local hi = opts.hi or #t + 1
-  local key = opts.key
+  local key_fn = make_key_fn(opts.key)
 
   if opts.bound == 'upper' then
-    return upper_bound(t, val, lo, hi, key)
+    return upper_bound(t, val, lo, hi, key_fn)
   else
-    return lower_bound(t, val, lo, hi, key)
+    return lower_bound(t, val, lo, hi, key_fn)
   end
 end
 

test/functional/lua/list_spec.lua

@@ -23,6 +23,8 @@ describe('vim.list', function()
         return x[1]
       end)
     )
+
+    eq({ { id = 1 }, { id = 2 } }, vim.list.unique({ { id = 1 }, { id = 2 }, { id = 1 } }, 'id'))
   end)
 
   --- Generate a list like { 1, 2, 2, 3, 3, 3, 4, 4, 4, 4, ...}.
@@ -52,6 +54,9 @@ describe('vim.list', function()
     local target = math.random(num)
     eq(vim.list.bisect(list, target, { bound = 'lower' }), index(target - 1) + 1)
     eq(vim.list.bisect(list, num + 1, { bound = 'lower' }), index(num) + 1)
+
+    local dict_list = { { value = 1 }, { value = 2 }, { value = 2 }, { value = 3 } }
+    eq(2, vim.list.bisect(dict_list, { value = 2 }, { key = 'value', bound = 'lower' }))
   end)
 
   it("vim.list.bisect(..., bound = { 'upper' })", function()
@@ -61,5 +66,8 @@ describe('vim.list', function()
     local target = math.random(num)
     eq(vim.list.bisect(list, target, { bound = 'upper' }), index(target) + 1)
     eq(vim.list.bisect(list, num + 1, { bound = 'upper' }), index(num) + 1)
+
+    local dict_list = { { value = 1 }, { value = 2 }, { value = 2 }, { value = 3 } }
+    eq(4, vim.list.bisect(dict_list, { value = 2 }, { key = 'value', bound = 'upper' }))
   end)
 end)