lsp_symbol picker: item in picker to respect the current or closest LSP symbol in main file?

[rrunner]

When I open the lsp_symbols picker, the default item in the picker is the first LSP symbol in the main file. How can I change the default item to be the current LSP symbol (if the cursor is on a symbol in the main file), or the closest LSP symbol (if the cursor is not on a symbol in the main file)?

[so11y]

me too

[drowning-cat]

Hide

---@param search_method? "cover_or_parent"|"cover_or_next"|"cover_or_prev"

---@module 'snacks'

-- Test keymap
vim.keymap.set("n", "<F1>", function()
  Snacks.picker.lsp_symbols()
end)

local follow_doc_symbol = function(picker, search_method)
  search_method = search_method or "cover_or_parent"

  -- NOTE: Items are assumed to be in order
  local list_items = vim.deepcopy(picker.finder.items)

  local win = picker:filter().current_win
  local buf = picker:filter().current_buf

  if vim.tbl_isempty(list_items) then
    return
  end

  local root = list_items[1].parent

  local set_target = function(item)
    picker.list:set_target(item.idx)
  end

  local get_lsp_range = function(item)
    -- return vim.range.lsp(buf, item.range, item.loc.encoding or "utf-16")
    return {
      start = { row = item.range.start.line, col = item.range.start.character },
      end_ = { row = item.range["end"].line, col = item.range["end"].character },
    }
  end
  local get_pos_range = function(pos)
    -- return vim.range(pos[1], pos[2], pos[1], pos[2] + 1, { buf = buf })
    return {
      start = { row = pos[1], col = pos[2] },
      end_ = { row = pos[1], col = pos[2] + 1 },
    }
  end
  local range_has = function(range1, range2)
    -- return range1:has(range2)
    local le = function(pos1, pos2)
      return pos1.row < pos2.row or (pos1.row == pos2.row and pos1.col <= pos2.col)
    end
    return le(range1.start, range2.start) and le(range2.end_, range1.end_)
  end
  local unwrap_range = function(range)
    return range.start.row, range.start.col, range.end_.row, range.end_.col
  end

  -- Add fields: +item.children, +item.vim_range
  for _, item in ipairs(list_items) do
    item.vim_range = get_lsp_range(item)
    item.children = item.children or {}
    item.parent.children = item.parent.children or {}
    table.insert(item.parent.children, item)
  end

  local _cursor = vim.api.nvim_win_get_cursor(win)
  local cursor_row, cursor_col = _cursor[1] - 1, _cursor[2]
  local cursor_range = get_pos_range({ cursor_row, cursor_col })

  local cursor_item

  local best_overlap = vim.iter(list_items):rev():find(function(item)
    return range_has(get_lsp_range(item), cursor_range)
  end)
  local best_same_line
  for _, item in ipairs(list_items) do
    local row, col = unwrap_range(item.vim_range)
    if row == cursor_row then
      if not best_same_line or col <= cursor_col then
        -- No break
        best_same_line = item
      end
    end
  end

  local rightmost_range = function(item1, item2)
    local row1, col1 = unwrap_range(item1.vim_range)
    local row2, col2 = unwrap_range(item2.vim_range)
    return (row1 < row2 or (row1 == row2 and col1 < col2)) and item2 or item1
  end

  if best_overlap and best_same_line then
    cursor_item = rightmost_range(best_same_line, best_overlap)
  else
    cursor_item = best_overlap or best_same_line
  end

  local children = cursor_item and cursor_item.children or root.children

  if cursor_item and (cursor_item == best_same_line or vim.tbl_isempty(children)) then
    set_target(cursor_item)
    return
  end

  local fallback

  if search_method == "cover_or_parent" then
    fallback = cursor_item
  end

  if search_method == "cover_or_next" then
    fallback = vim.iter(children):find(function(item)
      local row, col = unwrap_range(item.vim_range)
      return row > cursor_row or (row == cursor_row and col >= cursor_col)
    end)
  end

  if search_method == "cover_or_next" then
    fallback = vim.iter(children):rev():find(function(item)
      local _, _, end_row, end_col = unwrap_range(item.vim_range)
      return end_row < cursor_row or (end_row == cursor_row and end_col <= cursor_col)
    end)
  end

  if fallback then
    set_target(fallback)
    return
  end

  if not fallback and cursor_item then
    set_target(cursor_item)
    return
  end
end

local wait_items = function(picker, interval_ms, timeout_ms, cb)
  local deadline = vim.uv.now() + timeout_ms
  local function poll()
    if picker:count() > 0 then
      cb()
      return
    end
    if vim.uv.now() >= deadline then
      vim.notify("Timeout: No items found", vim.log.levels.WARN)
      return
    end
    vim.defer_fn(poll, interval_ms)
  end
  poll()
end

return {
  "folke/snacks.nvim",
  ---@type snacks.Config
  opts = {
    picker = {
      sources = {
        lsp_symbols = {
          on_show = function(picker)
            if picker.opts.workspace then ---@diagnostic disable-line
              return
            end
            wait_items(picker, 500, 10000, function()
              follow_doc_symbol(picker, "cover_or_parent")
            end)
          end,
        },
      },
    },
  },
}

[rrunner]

Thanks for the effort. I cannot verify this on my version of nvim (0.11.5) since vim.range is not available. It is probably a recently added module.

[drowning-cat]

Never mind, I updated the code with a "shimmed" vim.range (might be errors related to encoding).