fix: binding shortcuts to specific buffers (#238)

lua/avante/floating_window.lua

@@ -1,3 +1,5 @@
+local keymap = require("avante.utils.keymap")
+
 local api = vim.api
 
 local namespace = api.nvim_create_namespace("avante_floating_window")
@@ -237,19 +239,16 @@ function FloatingWindow:on(event, handler, options)
 end
 
 ---@param mode string|string[] check `:h :map-modes`
----@param key string|string[] key for the mapping
+---@param lhs string|string[]
 ---@param handler string | fun(): nil handler for the mapping
----@param opts? table<"'expr'"|"'noremap'"|"'nowait'"|"'remap'"|"'script'"|"'silent'"|"'unique'", boolean>
+---@param opts? vim.keymap.set.Opts
 ---@return nil
-function FloatingWindow:map(mode, key, handler, opts)
+function FloatingWindow:map(mode, lhs, handler, opts)
+  if not self.bufnr then
+    error("floating buffer not found.")
+  end
   local options = opts or {}
-  if type(key) == "string" then
+  keymap.set(self.bufnr, mode, lhs, handler, options)
-    vim.keymap.set(mode, key, handler, options)
-    return
-  end
-  for _, key_ in ipairs(key) do
-    vim.keymap.set(mode, key_, handler, options)
-  end
 end
 
 return FloatingWindow

lua/avante/sidebar.lua

@@ -1279,11 +1279,11 @@ function Sidebar:create_input()
     handle_submit(request)
   end
 
+  self.input:mount()
+
   self.input:map("n", Config.mappings.submit.normal, on_submit)
   self.input:map("i", Config.mappings.submit.insert, on_submit)
 
-  self.input:mount()
-
   api.nvim_set_option_value("filetype", "AvanteInput", { buf = self.input.bufnr })
 
   -- Setup completion
@@ -1463,6 +1463,8 @@ function Sidebar:render()
     },
   })
 
+  self.result:mount()
+
   self.result:on(event.BufWinEnter, function()
     xpcall(function()
       api.nvim_buf_set_name(self.result.bufnr, RESULT_BUF_NAME)
@@ -1479,8 +1481,6 @@ function Sidebar:render()
     self:close()
   end)
 
-  self.result:mount()
-
   self.input_container = Split({
     enter = false,
     relative = {

lua/avante/utils/buf_storage.lua

@@ -0,0 +1,34 @@
+-- Copy from: https://github.com/MunifTanjim/nui.nvim/blob/main/lua/nui/utils/buf_storage.lua
+local utils = require("avante.utils")
+
+local buf_storage = {
+  _registry = {},
+}
+
+---@param storage_name string
+---@param default_value any
+---@return table<number, any>
+function buf_storage.create(storage_name, default_value)
+  local storage = setmetatable({}, {
+    __index = function(tbl, bufnr)
+      rawset(tbl, bufnr, vim.deepcopy(utils.fallback(default_value, {})))
+
+      -- TODO: can `buf_storage.cleanup` be automatically (and reliably) triggered on `BufWipeout`?
+
+      return tbl[bufnr]
+    end,
+  })
+
+  buf_storage._registry[storage_name] = storage
+
+  return storage
+end
+
+---@param bufnr number
+function buf_storage.cleanup(bufnr)
+  for _, storage in pairs(buf_storage._registry) do
+    rawset(storage, bufnr, nil)
+  end
+end
+
+return buf_storage

lua/avante/utils/init.lua

@@ -343,4 +343,26 @@ function M.trim_spaces(s)
   return s:match("^%s*(.-)%s*$")
 end
 
+function M.fallback(v, default_value)
+  return type(v) == "nil" and default_value or v
+end
+
+-- luacheck: push no max comment line length
+---@param type_name "'nil'" | "'number'" | "'string'" | "'boolean'" | "'table'" | "'function'" | "'thread'" | "'userdata'" | "'list'" | '"map"'
+---@return boolean
+function M.is_type(type_name, v)
+  ---@diagnostic disable-next-line: deprecated
+  local islist = vim.islist or vim.tbl_islist
+  if type_name == "list" then
+    return islist(v)
+  end
+
+  if type_name == "map" then
+    return type(v) == "table" and not islist(v)
+  end
+
+  return type(v) == type_name
+end
+-- luacheck: pop
+
 return M

lua/avante/utils/keymap.lua

@@ -0,0 +1,103 @@
+-- Extracted from: https://github.com/MunifTanjim/nui.nvim/blob/main/lua/nui/utils/keymap.lua
+local buf_storage = require("avante.utils.buf_storage")
+local utils = require("avante.utils")
+
+local api = vim.api
+
+local keymap = {
+  storage = buf_storage.create("avante.utils.keymap", { _next_handler_id = 1, keys = {}, handlers = {} }),
+}
+
+---@param mode string
+---@param key string
+---@return string key_id
+local function get_key_id(mode, key)
+  return string.format("%s---%s", mode, vim.api.nvim_replace_termcodes(key, true, true, true))
+end
+
+---@param bufnr number
+---@param key_id string
+---@return integer|nil handler_id
+local function get_handler_id(bufnr, key_id)
+  return keymap.storage[bufnr].keys[key_id]
+end
+
+---@param bufnr number
+---@param mode string
+---@param key string
+---@param handler string|fun(): nil
+---@return { rhs: string, callback?: fun(): nil }|nil
+local function get_keymap_info(bufnr, mode, key, handler, overwrite)
+  local key_id = get_key_id(mode, key)
+
+  -- luacov: disable
+  if get_handler_id(bufnr, key_id) and not overwrite then
+    return nil
+  end
+  -- luacov: enable
+
+  local rhs, callback = "", nil
+
+  if type(handler) == "function" then
+    callback = handler
+  else
+    rhs = handler
+  end
+
+  return {
+    rhs = rhs,
+    callback = callback,
+  }
+end
+
+---@param bufnr number
+---@param mode string|string[]
+---@param lhs string|string[]
+---@param handler string|fun(): nil
+---@param opts? vim.keymap.set.Opts
+---@return nil
+function keymap.set(bufnr, mode, lhs, handler, opts, force)
+  if not utils.is_type("boolean", force) then
+    force = true
+  end
+
+  local keys = lhs
+  if type(lhs) ~= "table" then
+    keys = { lhs }
+  end
+  ---@cast keys -string
+
+  opts = opts or {}
+
+  if not utils.is_type("nil", opts.remap) then
+    opts.noremap = not opts.remap
+    opts.remap = nil
+  end
+
+  local modes = {}
+  if type(mode) == "string" then
+    modes = { mode }
+  else
+    modes = mode
+  end
+
+  for _, key in ipairs(keys) do
+    for _, mode_ in ipairs(modes) do
+      local keymap_info = get_keymap_info(bufnr, mode_, key, handler, force)
+      -- luacov: disable
+      if not keymap_info then
+        return false
+      end
+      -- luacov: enable
+
+      local options = vim.deepcopy(opts)
+      options.callback = keymap_info.callback
+
+      api.nvim_buf_set_keymap(bufnr, mode_, key, keymap_info.rhs, options)
+    end
+  end
+
+  return true
+end
+
+return keymap