From 36b23cef16c2c624c34bea213f01c06782d2ca40 Mon Sep 17 00:00:00 2001
From: teocns <59549574+teocns@users.noreply.github.com>
Date: Fri, 18 Oct 2024 00:41:40 -0600
Subject: [PATCH] feat: add repo map display (#727)

Co-authored-by: yetone <yetoneful@gmail.com>
---
 .github/workflows/ci.yaml |  2 +-
 lua/avante/config.lua     |  1 +
 lua/avante/init.lua       |  5 ++++
 lua/avante/repo_map.lua   | 58 ++++++++++++++++++++++++++++++++++++++-
 4 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index 8837a46..44a8644 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -61,7 +61,7 @@ jobs:
       - uses: Swatinem/rust-cache@v2
       - uses: dtolnay/rust-toolchain@master
         with:
-          toolchain: stable
+          toolchain: "1.80.1"
           components: clippy, rustfmt
       - name: Run rustfmt
         run: make rustlint
diff --git a/lua/avante/config.lua b/lua/avante/config.lua
index 8a4281b..da019ee 100644
--- a/lua/avante/config.lua
+++ b/lua/avante/config.lua
@@ -150,6 +150,7 @@ Respect and use existing conventions, libraries, etc that are already present in
       debug = "<leader>ad",
       hint = "<leader>ah",
       suggestion = "<leader>as",
+      repomap = "<leader>ar",
     },
     sidebar = {
       apply_all = "A",
diff --git a/lua/avante/init.lua b/lua/avante/init.lua
index eb7716d..c0d49d8 100644
--- a/lua/avante/init.lua
+++ b/lua/avante/init.lua
@@ -174,6 +174,11 @@ H.keymaps = function()
       function() M.toggle.suggestion() end,
       { desc = "avante: toggle suggestion" }
     )
+    Utils.safe_keymap_set("n", Config.mappings.toggle.repomap, function() require("avante.repo_map").show() end, {
+      desc = "avante: display repo map",
+      noremap = true,
+      silent = true,
+    })
   end
 
   if Config.behaviour.auto_suggestions then
diff --git a/lua/avante/repo_map.lua b/lua/avante/repo_map.lua
index 9e072e3..a0d2073 100644
--- a/lua/avante/repo_map.lua
+++ b/lua/avante/repo_map.lua
@@ -1,4 +1,6 @@
+local Popup = require("nui.popup")
 local Utils = require("avante.utils")
+local event = require("nui.utils.autocmd").event
 
 local filetype_map = {
   ["javascriptreact"] = "javascript",
@@ -38,6 +40,10 @@ function RepoMap._build_repo_map(project_root, file_ext)
   local filepaths = Utils.scan_directory(project_root, ignore_patterns, negate_patterns)
   vim.iter(filepaths):each(function(filepath)
     if not Utils.is_same_file_ext(file_ext, filepath) then return end
+    if not repo_map_lib then
+      Utils.error("Failed to load avante_repo_map")
+      return
+    end
     local definitions =
       repo_map_lib.stringify_definitions(RepoMap.get_ts_lang(filepath), Utils.file.read_content(filepath) or "")
     if definitions == "" then return end
@@ -141,7 +147,8 @@ function RepoMap._get_repo_map(file_ext)
   vim.api.nvim_create_autocmd({ "BufReadPost", "BufNewFile" }, {
     callback = function(ev)
       vim.defer_fn(function()
-        local filepath = vim.api.nvim_buf_get_name(ev.buf)
+        local ok, filepath = pcall(vim.api.nvim_buf_get_name, ev.buf)
+        if not ok or not filepath then return end
         if not vim.startswith(filepath, project_root) then return end
         local rel_filepath = Utils.relative_path(filepath)
         update_repo_map(rel_filepath)
@@ -152,4 +159,53 @@ function RepoMap._get_repo_map(file_ext)
   return repo_map
 end
 
+function RepoMap.show()
+  local file_ext = vim.fn.expand("%:e")
+  local repo_map = RepoMap.get_repo_map(file_ext)
+
+  if not repo_map or next(repo_map) == nil then
+    Utils.warn("The repo map is empty or not supported for this language: " .. file_ext)
+    return
+  end
+
+  local popup = Popup({
+    position = "50%",
+    enter = true,
+    focusable = true,
+    border = {
+      style = "rounded",
+      padding = { 1, 1 },
+      text = {
+        top = " Avante Repo Map ",
+        top_align = "center",
+      },
+    },
+    size = {
+      width = math.floor(vim.o.columns * 0.8),
+      height = math.floor(vim.o.lines * 0.8),
+    },
+  })
+
+  popup:mount()
+
+  popup:map("n", "q", function() popup:unmount() end, { noremap = true, silent = true })
+
+  popup:on(event.BufLeave, function() popup:unmount() end)
+
+  -- Format the repo map for display
+  local lines = {}
+  for _, entry in ipairs(repo_map) do
+    table.insert(lines, string.format("Path: %s", entry.path))
+    table.insert(lines, string.format("Lang: %s", entry.lang))
+    table.insert(lines, "Defs:")
+    for def_line in entry.defs:gmatch("[^\r\n]+") do
+      table.insert(lines, def_line)
+    end
+    table.insert(lines, "") -- Add an empty line between entries
+  end
+
+  -- Set the buffer content
+  vim.api.nvim_buf_set_lines(popup.bufnr, 0, -1, false, lines)
+end
+
 return RepoMap