Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Help] Adding treesitter highlighting for the entry_maker #3300

Closed
Ajaymamtora opened this issue Sep 21, 2024 · 13 comments
Closed

[Help] Adding treesitter highlighting for the entry_maker #3300

Ajaymamtora opened this issue Sep 21, 2024 · 13 comments
Labels
enhancement Enhancement to performance, inner workings or existent features

Comments

@Ajaymamtora
Copy link

Ajaymamtora commented Sep 21, 2024

Is your feature request related to a problem? Please describe.
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]

Sometimes its hard to understand the LSP results, so adding treesitter highlighting would make this much easier to understand. JetBrains IDEs have this functionality.

Describe the solution you'd like
A clear and concise description of what you want to happen.
Support highlighting the search result with treesitter

Additional context
Add any other context or screenshots about the feature request here.
I've gotten most of the way there I think but the highlights do not match the actual buffer so I think I'm doing something wrong.

Can someone help me or advise me where to go from here?

Heres the custom entry_maker so far:

local telescopeLspPickers = {}

local telescopeUtilities = require("telescope.utils")
local plenaryStrings = require("plenary.strings")
local telescopeEntryDisplayModule = require("telescope.pickers.entry_display")

-- Merged and case-insensitive icon map
local icon_map = require("core.ui.icons").lsp

-- Fixed widths for different parts of the entry
local ICON_WIDTH = 2
local FILENAME_WIDTH = 30
local LINE_COL_WIDTH = 8
local SEPARATOR_WIDTH = 1
-- The spacing between the segments up to the actual code segment
local PRE_CODE_SEGMENT_MARGINS = 3
local TEXT_OFFSET = ICON_WIDTH + FILENAME_WIDTH + LINE_COL_WIDTH + PRE_CODE_SEGMENT_MARGINS + (SEPARATOR_WIDTH * 3)

-- Cache for file highlights
local file_cache = {}

local function get_file_highlights(filename)
  if file_cache[filename] then
    return file_cache[filename]
  end

  local temp_bufnr = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_option(temp_bufnr, "buftype", "nofile")
  vim.api.nvim_buf_set_option(temp_bufnr, "bufhidden", "hide")
  vim.api.nvim_buf_set_option(temp_bufnr, "swapfile", false)

  local lines = vim.fn.readfile(filename)
  vim.api.nvim_buf_set_lines(temp_bufnr, 0, -1, false, lines)

  local lang = vim.filetype.match({ filename = filename })
  if lang then
    vim.api.nvim_buf_set_option(temp_bufnr, "filetype", lang)
  end

  local highlights = {}
  local parser = lang and vim.treesitter.get_parser(temp_bufnr, lang)
  local query = lang and vim.treesitter.query.get(lang, "highlights")

  if parser and query then
    local tree = parser:parse()[1]
    local root = tree:root()

    for i, line in ipairs(lines) do
      local line_highlights = {}
      for id, node, metadata in query:iter_captures(root, temp_bufnr, i - 1, i) do
        local name = query.captures[id]
        local start_row, start_col, end_row, end_col = node:range()
        if start_row == i - 1 and end_row == i - 1 then
          table.insert(line_highlights, {
            start_col = TEXT_OFFSET + start_col, -- Apply TEXT_OFFSET correctly
            end_col = TEXT_OFFSET + end_col, -- Apply TEXT_OFFSET correctly
            hl_group = name,
          })
        end
      end
      highlights[i] = line_highlights
    end
  end

  vim.api.nvim_buf_delete(temp_bufnr, { force = true })

  file_cache[filename] = {
    highlights = highlights,
    lines = lines,
  }

  return file_cache[filename]
end

function telescopeLspPickers.gen_from_lsp_symbols(opts)
  opts = opts or {}

  local displayer = telescopeEntryDisplayModule.create({
    separator = " ",
    items = {
      { width = ICON_WIDTH },
      { width = FILENAME_WIDTH },
      { width = LINE_COL_WIDTH },
      { remaining = true },
    },
  })

  return function(entry)
    local filename = entry.filename or vim.api.nvim_buf_get_name(entry.bufnr)
    local short_name = vim.fn.fnamemodify(filename, ":t")
    local lnum, col = entry.lnum, entry.col

    -- Retrieve symbol type, convert to lowercase and lookup the icon in the map
    local symbol_type = (entry.symbol_type or entry.type or "default"):lower()

    -- Use the icon_map with a fallback to a default icon
    local icon = icon_map[symbol_type] or "󰿰" -- Set a default icon if no match is found

    local file_data = get_file_highlights(filename)
    local text = file_data and file_data.lines[lnum] or entry.text or ""

    local display = function()
      local highlights = {}
      local truncated_name = plenaryStrings.truncate(short_name, FILENAME_WIDTH)
      local line_col = string.format("%4d:%-3d", lnum, col)
      local display_columns = {
        { icon },
        truncated_name,
        line_col,
        text,
      }

      if file_data and file_data.highlights and file_data.highlights[lnum] then
        for _, hl in ipairs(file_data.highlights[lnum]) do
          table.insert(highlights, {
            { hl.start_col, hl.end_col },
            hl.hl_group,
          })
        end
      end

      if entry.submatches then
        for _, match in ipairs(entry.submatches) do
          local s, e = match.start, match["end"]
          table.insert(highlights, { { TEXT_OFFSET + s, TEXT_OFFSET + e }, "TelescopeMatching" })
        end
      end

      return displayer(display_columns), highlights
    end

    return {
      value = entry,
      ordinal = short_name .. " " .. lnum .. ":" .. col .. " " .. text,
      display = display,
      filename = filename,
      lnum = lnum,
      col = col,
      text = text,
    }
  end
end

function telescopeLspPickers.cleanup()
  file_cache = {}
end

return telescopeLspPickers

And I'm using it via

-- ts opts:
{
  file_ignore_patterns = {},
  entry_maker = require("pickers.entry_makers.lsp").gen_from_lsp_references(),
}
Screenshot 2024-09-21 at 19 48 26
@Ajaymamtora Ajaymamtora added the enhancement Enhancement to performance, inner workings or existent features label Sep 21, 2024
@Ajaymamtora Ajaymamtora changed the title [Help] Add treesitter highlighting for the entry maker [Help] Adding treesitter highlighting for the entry_maker Sep 21, 2024
@jamestrew
Copy link
Contributor

Is that telescope results on the left and a regular buffer on the right?
If so, the difference in highlight is probably due to LSP semantic highlights.

@Ajaymamtora
Copy link
Author

Is that telescope results on the left and a regular buffer on the right?

If so, the difference in highlight is probably due to LSP semantic highlights.

Yes the left is the results

I thought that too but I disabled semantic highlights and the colour for this line specifically is the same.

I'll try again today to make sure I've disabled semantic highlights but I don't think that's the issue

@Ajaymamtora
Copy link
Author

Yeah I just had a look and I disabled the LSP completely and the highlights dont match the treesitter highlights.

@jamestrew
Copy link
Contributor

Maybe try :Inspect if you cursor over the COLOR_KEYS in the regular buffer. Could be other sources of highlight affecting it.

@Ajaymamtora
Copy link
Author

Maybe try :Inspect if you cursor over the COLOR_KEYS in the regular buffer. Could be other sources of highlight affecting it.

I inspected it and there's 3 items in the tree and the highest priority one is constant which is colouring it red. And the 1st one in the list matches up with the colour in the entry maker so maybe it's not resolving the higher priority highlight groups? I'm not sure if that's even a thing

@jamestrew
Copy link
Contributor

How does :Telescope current_buffer_fuzzy_find look with that buffer? It does result highlighting using treesitter so.
You code looks relatively similar so maybe you've already taken a look. Just wonder how current_buffer_fuzzy_find behaves.

@Ajaymamtora
Copy link
Author

Ajaymamtora commented Sep 22, 2024

I did try and look through the telescope project yeah, but I didn't realise there was something that did highlighting on entries, I could only find stuff for the previewer buffer.

Also current_buffer_fuzzy_find looks perfect, not sure what I'm doing differently, I'll have a look at how current_buffer_fuzzy_find does it.

Also the colour resolved to the middle one here this time:
image
so idk whats going on anymore with my code

@Ajaymamtora
Copy link
Author

I asked Claude to merge my picker with the current_buffer_fuzzy_find code and it fixed it,

local telescopeLspPickers = {}

local plenaryStrings = require("plenary.strings")
local telescopeEntryDisplayModule = require("telescope.pickers.entry_display")
local utils = require("telescope.utils")

-- Merged and case-insensitive icon map
local icon_map = require("core.ui.icons").lsp

-- Fixed widths for different parts of the entry
local ICON_WIDTH = 2
local FILENAME_WIDTH = 30
local LINE_COL_WIDTH = 8
local SEPARATOR_WIDTH = 1
-- The spacing between the segments up to the actual code segment
local PRE_CODE_SEGMENT_MARGINS = 3
local TEXT_OFFSET = ICON_WIDTH + FILENAME_WIDTH + LINE_COL_WIDTH + PRE_CODE_SEGMENT_MARGINS + (SEPARATOR_WIDTH * 3)

-- Cache for file highlights
local file_cache = {}

local function get_file_highlights(filename, opts)
  if file_cache[filename] then
    return file_cache[filename]
  end

  local temp_bufnr = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_option(temp_bufnr, "buftype", "nofile")
  vim.api.nvim_buf_set_option(temp_bufnr, "bufhidden", "hide")
  vim.api.nvim_buf_set_option(temp_bufnr, "swapfile", false)

  local lines = vim.fn.readfile(filename)
  vim.api.nvim_buf_set_lines(temp_bufnr, 0, -1, false, lines)

  local filetype = vim.filetype.match({ filename = filename })
  if filetype then
    vim.api.nvim_buf_set_option(temp_bufnr, "filetype", filetype)
  end

  local highlights = {}
  local lang = vim.treesitter.language.get_lang(filetype) or filetype

  if opts.results_ts_highlight and lang and utils.has_ts_parser(lang) then
    local parser = vim.treesitter.get_parser(temp_bufnr, lang)
    local query = vim.treesitter.query.get(lang, "highlights")
    local root = parser:parse()[1]:root()

    for id, node in query:iter_captures(root, temp_bufnr, 0, -1) do
      local hl = "@" .. query.captures[id]
      if hl and type(hl) ~= "number" then
        local row1, col1, row2, col2 = node:range()
        if row1 == row2 then
          local row = row1 + 1
          if not highlights[row] then
            highlights[row] = {}
          end
          for index = col1, col2 - 1 do
            highlights[row][index] = hl
          end
        else
          local row = row1 + 1
          if not highlights[row] then
            highlights[row] = {}
          end
          for index = col1, #lines[row] - 1 do
            highlights[row][index] = hl
          end
          while row < row2 do
            row = row + 1
            if not highlights[row] then
              highlights[row] = {}
            end
            for index = 0, #(lines[row] or {}) - 1 do
              highlights[row][index] = hl
            end
          end
        end
      end
    end
  end

  vim.api.nvim_buf_delete(temp_bufnr, { force = true })

  file_cache[filename] = {
    highlights = highlights,
    lines = lines,
  }

  return file_cache[filename]
end

function telescopeLspPickers.gen_from_lsp_symbols(opts)
  opts = opts or {}
  opts.results_ts_highlight = vim.F.if_nil(opts.results_ts_highlight, true)

  local displayer = telescopeEntryDisplayModule.create({
    separator = " ",
    items = {
      { width = ICON_WIDTH },
      { width = FILENAME_WIDTH },
      { width = LINE_COL_WIDTH },
      { remaining = true },
    },
  })

  return function(entry)
    local filename = entry.filename or vim.api.nvim_buf_get_name(entry.bufnr)
    local short_name = vim.fn.fnamemodify(filename, ":t")
    local lnum, col = entry.lnum, entry.col

    local symbol_type = (entry.symbol_type or entry.type or "default"):lower()
    local icon = icon_map[symbol_type] or "󰿰"

    local file_data = get_file_highlights(filename, opts)
    local text = file_data and file_data.lines[lnum] or entry.text or ""

    local display = function()
      local highlights = {}
      local truncated_name = plenaryStrings.truncate(short_name, FILENAME_WIDTH)
      local display_columns = {
        string.format("%-" .. ICON_WIDTH .. "s", icon),
        string.format("%-" .. FILENAME_WIDTH .. "s", truncated_name),
        string.format("%4d:%-3d", lnum, col),
        text,
      }

      if file_data and file_data.highlights and file_data.highlights[lnum] then
        for col, hl_group in pairs(file_data.highlights[lnum]) do
          table.insert(highlights, {
            { TEXT_OFFSET + col, TEXT_OFFSET + col + 1 },
            hl_group,
          })
        end
      end

      if entry.submatches then
        for _, match in ipairs(entry.submatches) do
          local s, e = match.start, match["end"]
          table.insert(highlights, { { TEXT_OFFSET + s, TEXT_OFFSET + e }, "TelescopeMatching" })
        end
      end

      return displayer(display_columns), highlights
    end

    return {
      value = entry,
      ordinal = short_name .. " " .. lnum .. ":" .. col .. " " .. text,
      display = display,
      filename = filename,
      lnum = lnum,
      col = col,
      text = text,
    }
  end
end

function telescopeLspPickers.cleanup()
  file_cache = {}
end

return telescopeLspPickers

and it works great for LSP pickers, so thanks for the help!

@Ajaymamtora
Copy link
Author

Although when listing document symbols the highlighting is off by 1 char, and adjusting the offset doesnt work. Im using Aerial telescope, not sure how that makes any sense

@Ajaymamtora
Copy link
Author

How does :Telescope current_buffer_fuzzy_find look with that buffer? It does result highlighting using treesitter so. You code looks relatively similar so maybe you've already taken a look. Just wonder how current_buffer_fuzzy_find behaves.

Ive realised my code has a bug, when the filename is too long, the highlights go out of sync, even though theyre truncated, I cant figure this out.

Can you help me please?

image

local telescopeLspPickers = {}

local plenaryStrings = require("plenary.strings")
local telescopeEntryDisplayModule = require("telescope.pickers.entry_display")

-- Add a debug function
local function debug_print(...)
  print(string.format(...))
end

-- Merged and case-insensitive icon map
local icon_map = require("core.ui.icons").lsp

-- Fixed widths for different parts of the entry
local ICON_WIDTH = 2
local FILENAME_WIDTH = 30
local LINE_COL_WIDTH = 8
local SEPARATOR_WIDTH = 1
local TS_PADDING = 3

-- Calculate the fixed highlight start position
local HIGHLIGHT_START = ICON_WIDTH + FILENAME_WIDTH + LINE_COL_WIDTH + TS_PADDING + (3 * SEPARATOR_WIDTH)

-- Cache for file highlights
local file_cache = {}

local function detect_filetype_safe(filepath)
  local current_buf = vim.api.nvim_get_current_buf()
  local current_win = vim.api.nvim_get_current_win()

  -- Create a new hidden buffer
  local buf = vim.api.nvim_create_buf(false, true)

  -- Read file content
  local lines = vim.fn.readfile(filepath)

  -- Set buffer content
  vim.api.nvim_buf_set_lines(buf, 0, -1, false, lines)

  -- Use vim.filetype.match to detect filetype
  local filetype = vim.filetype.match({ buf = buf, filename = filepath })

  -- Clean up
  vim.api.nvim_buf_delete(buf, { force = true })
  vim.api.nvim_set_current_win(current_win)
  vim.api.nvim_set_current_buf(current_buf)

  return filetype or ""
end

local function collect_buf_highlights(bufnr, lines)
  local parser = vim.treesitter.get_parser(bufnr)
  if not parser then
    return {}
  end

  local lang = parser:lang()
  local root = parser:parse()[1]:root()

  local highlights = {}
  local query = vim.treesitter.query.get(lang, "highlights")
  if query then
    for id, node, metadata in query:iter_captures(root, bufnr, 0, -1) do
      local hl = "@" .. query.captures[id]
      if hl and type(hl) ~= "number" then
        local start_row, start_col, end_row, end_col = node:range()
        for row = start_row, end_row do
          highlights[row + 1] = highlights[row + 1] or {}
          local col_start = row == start_row and start_col or 0
          local col_end = row == end_row and end_col or #lines[row + 1]
          for col = col_start, col_end - 1 do
            highlights[row + 1][col] = hl
          end
        end
      end
    end
  end
  return highlights
end

local function get_file_highlights(filename, opts)
  if file_cache[filename] then
    return file_cache[filename]
  end

  local lines = vim.fn.readfile(filename)
  local filetype = detect_filetype_safe(filename)

  local bufnr = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)
  vim.api.nvim_buf_set_option(bufnr, "filetype", filetype)

  local highlights = collect_buf_highlights(bufnr, lines)

  vim.api.nvim_buf_delete(bufnr, { force = true })

  file_cache[filename] = {
    highlights = highlights,
    lines = lines,
    filetype = filetype,
  }

  return file_cache[filename]
end

-- New function to ensure filename is exactly FILENAME_WIDTH characters
local function format_filename(filename)
  if #filename <= FILENAME_WIDTH then
    return filename .. string.rep(" ", FILENAME_WIDTH - #filename)
  else
    return plenaryStrings.truncate(filename, FILENAME_WIDTH - 1) .. ""
  end
end

function telescopeLspPickers.gen_from_lsp_symbols(opts)
  opts = opts or {}
  opts.results_ts_highlight = vim.F.if_nil(opts.results_ts_highlight, true)

  return function(entry)
    local filename = entry.filename or vim.api.nvim_buf_get_name(entry.bufnr)
    local short_name = vim.fn.fnamemodify(filename, ":t")
    local lnum, col = entry.lnum, entry.col
    local has_line_col = lnum ~= nil and col ~= nil

    local symbol_type = (entry.symbol_type or entry.type or "default"):lower()
    local icon = icon_map[symbol_type] or "󰿰"

    local file_data = get_file_highlights(filename, opts)
    local original_text = file_data and file_data.lines[lnum] or entry.text or ""
    local text = original_text -- Removed vim.trim

    local displayer_items = {
      { width = ICON_WIDTH },
      { width = FILENAME_WIDTH },
      { width = LINE_COL_WIDTH },
      { remaining = true },
    }

    local displayer = telescopeEntryDisplayModule.create({
      separator = "-",
      items = displayer_items,
    })

    local display = function()
      debug_print("Entering display function")
      local highlights = {}
      local formatted_name = format_filename(short_name)
      debug_print("Formatted filename: %s", formatted_name)

      local display_columns = {
        icon,
        formatted_name,
        has_line_col and string.format("%4d:%-3d", lnum, col) or string.rep(" ", LINE_COL_WIDTH),
        text,
      }
      debug_print(
        "Display columns: icon=%s, filename=%s, line_col=%s, text=%s",
        icon,
        formatted_name,
        display_columns[3],
        text
      )

      if file_data and file_data.highlights and file_data.highlights[lnum] then
        debug_print("Highlights found for line %d", lnum)
        for hlcol, hl_group in pairs(file_data.highlights[lnum]) do
          if hlcol >= 0 and hlcol < #text then
            local hl_start = HIGHLIGHT_START + hlcol
            local hl_end = hl_start + 1
            table.insert(highlights, {
              { hl_start, hl_end },
              hl_group,
            })
            debug_print("Added highlight: col=%d, start=%d, end=%d, group=%s", hlcol, hl_start, hl_end, hl_group)
          end
        end
      else
        debug_print("No highlights found for line %d", lnum)
      end

      debug_print("Number of highlights added: %d", #highlights)
      local result, result_highlights = displayer(display_columns), highlights
      debug_print("Display function completed")
      return result, result_highlights
    end

    local ordinal = short_name
    if has_line_col then
      ordinal = ordinal .. " " .. lnum .. ":" .. col
    end
    ordinal = ordinal .. " " .. text

    return {
      value = entry,
      ordinal = ordinal,
      display = display,
      filename = filename,
      lnum = lnum,
      col = col,
      text = text,
    }
  end
end

function telescopeLspPickers.cleanup()
  file_cache = {}
end

return telescopeLspPickers

Is there a way to make the string a fixed length using require("telescope.pickers.entry_display").create?

@Ajaymamtora Ajaymamtora reopened this Sep 23, 2024
@Ajaymamtora
Copy link
Author

Fixed by defining a custom truncate function

local telescopeLspPickers = {}

local plenaryStrings = require("plenary.strings")

-- Add a debug function
local function debug_print(...)
  print(string.format(...))
end

-- Merged and case-insensitive icon map
local icon_map = require("core.ui.icons").lsp

-- Fixed widths for different parts of the entry
local ICON_WIDTH = 2
local FILENAME_WIDTH = 30
local LINE_COL_WIDTH = 8
local SEPARATOR_WIDTH = 1
local TS_PADDING = 3

-- Calculate the fixed highlight start position
local HIGHLIGHT_START = ICON_WIDTH + FILENAME_WIDTH + LINE_COL_WIDTH + TS_PADDING + (3 * SEPARATOR_WIDTH)

-- Cache for file highlights
local file_cache = setmetatable({}, {
  __index = function(t, k)
    local v = setmetatable({}, { __mode = "v" })
    rawset(t, k, v)
    return v
  end,
})

local function detect_filetype_safe(filepath)
  local current_buf = vim.api.nvim_get_current_buf()
  local current_win = vim.api.nvim_get_current_win()

  -- Create a new hidden buffer
  local buf = vim.api.nvim_create_buf(false, true)

  -- Read file content
  local lines = vim.fn.readfile(filepath)

  -- Set buffer content
  vim.api.nvim_buf_set_lines(buf, 0, -1, false, lines)

  -- Use vim.filetype.match to detect filetype
  local filetype = vim.filetype.match({ buf = buf, filename = filepath })

  -- Clean up
  vim.api.nvim_buf_delete(buf, { force = true })
  vim.api.nvim_set_current_win(current_win)
  vim.api.nvim_set_current_buf(current_buf)

  return filetype or ""
end

local function collect_buf_highlights(bufnr, lines)
  local parser = vim.treesitter.get_parser(bufnr)
  if not parser then
    return {}
  end

  local lang = parser:lang()
  local tree = parser:parse()[1]
  if not tree then
    return {}
  end
  local root = tree:root()

  local highlights = {}
  local query = vim.treesitter.query.get(lang, "highlights")
  if query then
    for id, node, _ in query:iter_captures(root, bufnr, 0, -1) do
      local hl = "@" .. query.captures[id]
      if hl and type(hl) ~= "number" then
        local start_row, start_col, end_row, end_col = node:range()
        for row = start_row, end_row do
          highlights[row + 1] = highlights[row + 1] or {}
          local col_start = row == start_row and start_col or 0
          local col_end = row == end_row and end_col or #(lines[row + 1] or "")
          for col = col_start, col_end - 1 do
            highlights[row + 1][col] = hl
          end
        end
      end
    end
  end
  return highlights
end

local function get_file_highlights(filename, _)
  if file_cache[filename].highlights then
    return file_cache[filename]
  end

  local lines = vim.fn.readfile(filename)
  local filetype = detect_filetype_safe(filename)

  local bufnr = vim.api.nvim_create_buf(false, true)
  vim.api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)
  vim.api.nvim_buf_set_option(bufnr, "filetype", filetype)

  local highlights = collect_buf_highlights(bufnr, lines)

  vim.api.nvim_buf_delete(bufnr, { force = true })

  file_cache[filename] = {
    highlights = highlights,
    lines = lines,
    filetype = filetype,
  }

  return file_cache[filename]
end

-- Custom truncate function to keep track of truncation offset
local function custom_truncate(str, max_len, is_icon)
  if is_icon or #str <= max_len then
    return str, false, 0
  end
  return str:sub(1, max_len - 3) .. "...", true, 2
end

local function entry_display_create(configuration)
  local state = require("telescope.state")
  local resolve = require("telescope.config.resolve")

  local generator = {}
  for i, v in ipairs(configuration.items) do
    if v.width then
      local justify = v.right_justify
      local width
      table.insert(generator, function(item)
        if width == nil then
          local status = state.get_status(vim.F.if_nil(configuration.prompt_bufnr, vim.api.nvim_get_current_buf()))
          local s = {}
          s[1] = vim.api.nvim_win_get_width(status.layout.results.winid) - #status.picker.selection_caret
          s[2] = vim.api.nvim_win_get_height(status.layout.results.winid)
          width = resolve.resolve_width(v.width)(nil, s[1], s[2])
        end
        local is_icon = i == 1 -- Check if this is the icon column
        local truncated, was_truncated, offset = custom_truncate(tostring(item), width, is_icon)
        return plenaryStrings.align_str(truncated, width, justify), nil, was_truncated, #truncated, offset
      end)
    else
      table.insert(generator, function(item)
        return tostring(item), nil, false, #tostring(item), 0
      end)
    end
  end

  return function(self, picker)
    local results = {}
    local highlights = {}
    local truncation_info = {}
    local cumulative_length = 0
    for i, _ in ipairs(configuration.items) do
      if self[i] ~= nil then
        local str, _, was_truncated, displayed_length, offset = generator[i](self[i], picker)
        truncation_info[i] = {
          original = tostring(self[i]),
          displayed = str,
          was_truncated = was_truncated,
          displayed_length = displayed_length,
          offset = offset,
        }
        table.insert(results, str)
        cumulative_length = cumulative_length + #str + (#configuration.separator or 1)
      end
    end

    local final_str = table.concat(results, configuration.separator or "")
    return final_str, highlights, truncation_info
  end
end

function telescopeLspPickers.gen_from_lsp_symbols(opts)
  opts = opts or {}
  opts.results_ts_highlight = vim.F.if_nil(opts.results_ts_highlight, true)

  return function(entry)
    local filename = entry.filename or vim.api.nvim_buf_get_name(entry.bufnr)
    local short_name = vim.fn.fnamemodify(filename, ":t")
    local lnum, col = entry.lnum, entry.col
    local has_line_col = lnum ~= nil and col ~= nil

    local symbol_type = (entry.symbol_type or entry.type or "default"):lower()
    local icon = icon_map[symbol_type] or "󰿰"

    local file_data = get_file_highlights(filename, opts)
    local original_text = file_data and file_data.lines[lnum] or entry.text or ""
    local text = original_text -- Removed vim.trim

    local displayer_items = {
      { width = ICON_WIDTH },
      { width = FILENAME_WIDTH },
      { width = LINE_COL_WIDTH },
      { remaining = true },
    }

    local displayer = entry_display_create({
      separator = " ",
      items = displayer_items,
    })

    local display = function()
      debug_print("Entering display function")
      local highlights = {}
      local formatted_name = plenaryStrings.truncate(short_name, FILENAME_WIDTH)
      debug_print("Formatted filename: %s", formatted_name)

      local display_columns = {
        icon,
        formatted_name,
        has_line_col and string.format("%4d:%-3d", lnum, col) or string.rep(" ", LINE_COL_WIDTH),
        text,
      }
      debug_print(
        "Display columns: icon=%s, filename=%s, line_col=%s, text=%s",
        icon,
        formatted_name,
        display_columns[3],
        text
      )

      local result, _, truncation_info = displayer(display_columns)

      if file_data and file_data.highlights and file_data.highlights[lnum] then
        debug_print("Highlights found for line %d", lnum)
        local text_info = truncation_info[4]
        local is_truncated = text_info.was_truncated
        local displayed_length = text_info.displayed_length
        local truncation_offset = text_info.offset

        debug_print(
          "Is truncated: %s, Displayed length: %d, Truncation offset: %d",
          tostring(is_truncated),
          displayed_length,
          truncation_offset
        )

        for hlcol, hl_group in pairs(file_data.highlights[lnum]) do
          if hlcol >= 0 and hlcol < displayed_length - truncation_offset then
            local hl_start = HIGHLIGHT_START + hlcol
            local hl_end = hl_start + 1
            table.insert(highlights, {
              { hl_start, hl_end },
              hl_group,
            })
            debug_print("Added highlight: col=%d, start=%d, end=%d, group=%s", hlcol, hl_start, hl_end, hl_group)
          end
        end
      else
        debug_print("No highlights found for line %d", lnum)
      end

      debug_print("Number of highlights added: %d", #highlights)
      return result, highlights
    end

    local ordinal = short_name
    if has_line_col then
      ordinal = ordinal .. " " .. lnum .. ":" .. col
    end
    ordinal = ordinal .. " " .. text

    return {
      value = entry,
      ordinal = ordinal,
      display = display,
      filename = filename,
      lnum = lnum,
      col = col,
      text = text,
    }
  end
end

function telescopeLspPickers.cleanup()
  file_cache = setmetatable({}, {
    __index = function(t, k)
      local v = setmetatable({}, { __mode = "v" })
      rawset(t, k, v)
      return v
    end,
  })
end

return telescopeLspPickers

@fdschmidt93
Copy link
Member

You might wanna look at this if performance is a concern

fdschmidt93/telescope-egrepify.nvim#50

@Ajaymamtora
Copy link
Author

You might wanna look at this if performance is a concern

fdschmidt93/telescope-egrepify.nvim#50

Cheers, performance seems to be fine but I'll make use of that anyway 👍🏼

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement Enhancement to performance, inner workings or existent features
Projects
None yet
Development

No branches or pull requests

3 participants