Repository: neovim/neovim
File: runtime/lua/vim/shared.lua (Lua)

Code before:
---
--- i = vim.list.lower_bound(t, 4)
--- -- i is now 7
--- ```
---
--- or you can use it together with |vim.list.upper_bound()|
--- to obtain the range of occurrences.
---@see |vim.list.upper_bound()|
---@generic T
---@param t T[] A comparable list.
---@param val T The value to search.
---@param opts? vim.list.bisect.Opts
---@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.
function vim.list.lower_bound(t, val, opts)
  vim.validate('t', t, 'table')
  vim.validate('opts', opts, 'table', true)

  opts = opts or {}
  local lo = opts.lo or 1
  local hi = opts.hi or #t + 1

  while lo < hi do
    local mid = bit.rshift(lo + hi, 1) -- Equivalent to floor((lo + hi) / 2).
    if key_fn(t[mid], opts.key) < key_fn(val, opts.key) then
      lo = mid + 1
    else
      hi = mid
    end
  end

  return lo
end

--- Do a upper-bound binary search of the list in the half-open range [lo, hi),
--- returns the index in {t} after the last occurrence of the given {val}.
---
--- Example:
--- ```lua
---
--- local t = { 1, 2, 2, 3, 3, 3 }
--- local start = vim.list.lower_bound(t, 3)
--- -- start is now 4
--- local end = vim.list.upper_bound(t, 2)
--- -- end is now 7
---
--- -- 3 is in [start, end)
--- for i = start, end - 1 do
---   print(t[i])
--- end

Review comment (suggestion):
```suggestion
  local val_key = key_fn(val, key)
  
  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
```

Code after:
---@param t T[]
---@param val T
---@param key? fun(val: any): any
---@param lo integer
---@param hi integer
---@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 bit = require('bit') -- Load bitop on demand
  local val_key = key_fn(val, key)
  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
      lo = mid + 1
    else
      hi = mid
    end
  end
  return lo
end

---@generic T
---@param t T[]
---@param val T
---@param key? fun(val: any): any
---@param lo integer
---@param hi integer
---@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 bit = require('bit') -- Load bitop on demand
  local val_key = key_fn(val, key)
  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
      hi = mid
    else
      lo = mid + 1
    end
  end
  return lo
end

--- Search for a position in a sorted list {t}
--- where {val} can be inserted while keeping the list sorted.
---
--- Use {bound} to determine whether to return the first or the last position,
--- defaults to "lower", i.e., the first position.
---