fix: synchronize the size of the floating window and split window (ye…

…tone#151)
SiyuLiu0329 · Aug 22, 2024 · 67c0be2 · 67c0be2
1 parent 8c8d864
commit 67c0be2
Show file tree

Hide file tree

Showing 3 changed files with 183 additions and 31 deletions.
diff --git a/lua/avante/floating_window.lua b/lua/avante/floating_window.lua
@@ -12,6 +12,9 @@ api.nvim_set_hl(namespace, "FloatBorder", { link = "Normal" })
 ---@field buf_options table | nil
 ---@field win_options table | nil
 ---@field float_options table | nil
+---@field on_mount_handlers table | nil
+---@field on_unmount_handlers table | nil
+---@field augroup integer | nil
 local FloatingWindow = {}
 FloatingWindow.__index = FloatingWindow
 
@@ -37,9 +40,128 @@ function FloatingWindow.new(opts)
   instance.buf_options = opts.buf_options or {}
   instance.win_options = opts.win_options or {}
   instance.float_options = opts.float_options or {}
+  instance.on_mount_handlers = {}
+  instance.on_unmount_handlers = {}
+  instance.augroup = nil
   return instance
 end
 
+---@param split_winid integer
+---@param opts opts
+---@return FloatingWindow
+function FloatingWindow.from_split_win(split_winid, opts)
+  local split_win_width = api.nvim_win_get_width(split_winid)
+  local split_win_height = api.nvim_win_get_height(split_winid)
+
+  local calc_floating_win_size = function(width, height)
+    return {
+      width = math.max(width - 2, 0),
+      height = math.max(height - 3, 0),
+    }
+  end
+
+  local floating_win_size = calc_floating_win_size(split_win_width, split_win_height)
+
+  local float_opts_ = vim.tbl_deep_extend("force", {
+    relative = "win",
+    win = split_winid,
+    width = floating_win_size.width,
+    height = floating_win_size.height,
+    row = 1,
+    col = 0,
+    style = "minimal",
+    border = { " ", " ", " ", " ", " ", " ", " ", " " },
+  }, opts.float_options or {})
+
+  local win_opts_ = vim.tbl_deep_extend("force", {}, opts.win_options or {})
+
+  local buf_opts_ = vim.tbl_deep_extend("force", {}, opts.buf_options or {})
+
+  local floating_win = FloatingWindow({
+    buf_options = buf_opts_,
+    win_options = win_opts_,
+    float_options = float_opts_,
+  })
+
+  floating_win:on_mount(function(winid)
+    api.nvim_create_autocmd("WinResized", {
+      group = floating_win.augroup,
+      callback = function()
+        if not api.nvim_win_is_valid(split_winid) or not api.nvim_win_is_valid(winid) then
+          return
+        end
+
+        local current_width = api.nvim_win_get_width(winid)
+        local current_height = api.nvim_win_get_height(winid)
+
+        if current_width == floating_win_size.width and current_height == floating_win_size.height then
+          return
+        end
+
+        floating_win_size.width = current_width
+        floating_win_size.height = current_height
+
+        api.nvim_win_set_height(split_winid, current_height + 3)
+        api.nvim_win_set_width(split_winid, current_width + 2)
+      end,
+    })
+
+    api.nvim_create_autocmd("WinResized", {
+      group = floating_win.augroup,
+      callback = function()
+        if not api.nvim_win_is_valid(split_winid) or not api.nvim_win_is_valid(winid) then
+          return
+        end
+
+        local current_split_win_width = api.nvim_win_get_width(split_winid)
+        local current_split_win_height = api.nvim_win_get_height(split_winid)
+
+        if current_split_win_width == split_win_width and current_split_win_height == split_win_height then
+          return
+        end
+
+        split_win_width = current_split_win_width
+        split_win_height = current_split_win_height
+
+        local current_floating_win_size = calc_floating_win_size(current_split_win_width, current_split_win_height)
+
+        local old_floating_win_options = api.nvim_win_get_config(winid)
+        local new_floating_win_options = vim.tbl_deep_extend("force", old_floating_win_options, {
+          width = current_floating_win_size.width,
+          height = current_floating_win_size.height,
+        })
+        api.nvim_win_set_config(winid, new_floating_win_options)
+      end,
+    })
+  end)
+
+  return floating_win
+end
+
+function FloatingWindow:__gc()
+  self:unmount()
+end
+
+function FloatingWindow:__tostring()
+  return "FloatingWindow"
+end
+
+function FloatingWindow:__eq(other)
+  return self.winid == other.winid
+end
+
+---@param handler fun(number, number): nil
+---@return nil
+function FloatingWindow:on_mount(handler)
+  table.insert(self.on_mount_handlers, handler)
+end
+
+---@param handler fun(nuber, nubmer): nil
+---@return nil
+function FloatingWindow:on_unmount(handler)
+  table.insert(self.on_unmount_handlers, handler)
+end
+
 ---@return nil
 function FloatingWindow:mount()
   self.bufnr = api.nvim_create_buf(false, true)
@@ -50,15 +172,30 @@ function FloatingWindow:mount()
 
   self.winid = api.nvim_open_win(self.bufnr, self.enter, self.float_options)
 
+  self.augroup = api.nvim_create_augroup("avante_floating_window_" .. tostring(self.winid), { clear = true })
+
   for option, value in pairs(self.win_options) do
     api.nvim_set_option_value(option, value, { win = self.winid })
   end
 
   api.nvim_win_set_hl_ns(self.winid, namespace)
+
+  for _, handler in ipairs(self.on_mount_handlers) do
+    handler(self.winid, self.bufnr)
+  end
 end
 
 ---@return nil
 function FloatingWindow:unmount()
+  for _, handler in ipairs(self.on_unmount_handlers) do
+    handler(self.winid, self.bufnr)
+  end
+
+  if self.augroup ~= nil then
+    pcall(api.nvim_delete_augroup, self.augroup)
+    self.augroup = nil
+  end
+
   if self.bufnr and api.nvim_buf_is_valid(self.bufnr) then
     api.nvim_buf_delete(self.bufnr, { force = true })
     self.bufnr = nil

diff --git a/lua/avante/sidebar.lua b/lua/avante/sidebar.lua
@@ -31,12 +31,12 @@ local Sidebar = {}
 ---@field augroup integer
 ---@field code avante.CodeState
 ---@field winids table<string, integer> this table stores the winids of the sidebar components (result_container, result, selected_code_container, selected_code, input_container, input), even though they are destroyed.
----@field result_container NuiSplit | nil
----@field result NuiSplit | nil
----@field selected_code_container NuiSplit | nil
----@field selected_code NuiSplit | nil
----@field input_container NuiSplit | nil
----@field input NuiSplit | nil
+---@field result_container AvanteComp | nil
+---@field result FloatingWindow | nil
+---@field selected_code_container AvanteComp | nil
+---@field selected_code FloatingWindow | nil
+---@field input_container AvanteComp | nil
+---@field input FloatingWindow | nil
 
 ---@param id integer the tabpage id retrieved from api.nvim_get_current_tabpage()
 function Sidebar:new(id)
@@ -580,6 +580,37 @@ function Sidebar:on_mount()
     end,
   })
 
+  local input_container_win_width = api.nvim_win_get_width(self.input_container.winid)
+  local input_container_win_height = api.nvim_win_get_height(self.input_container.winid)
+
+  api.nvim_create_autocmd("WinResized", {
+    group = self.augroup,
+    callback = function()
+      if not api.nvim_win_is_valid(self.input_container.winid) or not api.nvim_win_is_valid(self.input.winid) then
+        return
+      end
+
+      local current_input_container_win_width = api.nvim_win_get_width(self.input_container.winid)
+      local current_input_container_win_height = api.nvim_win_get_height(self.input_container.winid)
+
+      if
+        current_input_container_win_width == input_container_win_width
+        and current_input_container_win_height == input_container_win_height
+      then
+        return
+      end
+
+      input_container_win_width = current_input_container_win_width
+      input_container_win_height = current_input_container_win_height
+
+      local old_floating_win_options = api.nvim_win_get_config(self.input.winid)
+      local new_floating_win_options = vim.tbl_deep_extend("force", old_floating_win_options, {
+        width = current_input_container_win_width - 2,
+      })
+      api.nvim_win_set_config(self.input.winid, new_floating_win_options)
+    end,
+  })
+
   api.nvim_create_autocmd("WinClosed", {
     group = self.augroup,
     callback = function(args)
@@ -1141,30 +1172,14 @@ function Sidebar:get_selected_code_size()
 end
 
 local function create_floating_window_for_split(split_winid, buf_opts, win_opts, float_opts)
-  local height = api.nvim_win_get_height(split_winid)
-  local width = api.nvim_win_get_width(split_winid)
-
-  local float_opts_ = vim.tbl_deep_extend("force", {
-    relative = "win",
-    win = split_winid,
-    width = math.max(width - 2, 0),
-    height = math.max(height - 3, 0),
-    row = 1,
-    col = 0,
-    style = "minimal",
-    border = { " ", " ", " ", " ", " ", " ", " ", " " },
-  }, float_opts or {})
-
-  local win_opts_ = vim.tbl_deep_extend("force", get_win_options(), {
-    winhighlight = "NormalFloat:Normal,FloatBorder:Normal",
-  }, win_opts or {})
+  local win_opts_ = vim.tbl_deep_extend("force", get_win_options(), win_opts or {})
 
   local buf_opts_ = vim.tbl_deep_extend("force", buf_options, buf_opts or {})
 
-  local floating_win = FloatingWindow({
+  local floating_win = FloatingWindow.from_split_win(split_winid, {
     buf_options = buf_opts_,
     win_options = win_opts_,
-    float_options = float_opts_,
+    float_options = float_opts,
   })
 
   return floating_win

diff --git a/lua/avante/types.lua b/lua/avante/types.lua
@@ -1,26 +1,26 @@
 ---@meta
 
----@class NuiSplit
+---@class AvanteComp
 ---@field winid integer | nil
 ---@field bufnr integer | nil
-local AvanteSplit = require("nui.split")
+local AvanteComp = {}
 
 ---@return nil
-function AvanteSplit:mount() end
+function AvanteComp:mount() end
 
 ---@return nil
-function AvanteSplit:unmount() end
+function AvanteComp:unmount() end
 
 ---@param event string | string[]
 ---@param handler string | function
 ---@param options? table<"'once'" | "'nested'", boolean>
 ---@return nil
-function AvanteSplit:on(event, handler, options) end
+function AvanteComp:on(event, handler, options) end
 
 -- set keymap for this split
 ---@param mode string check `:h :map-modes`
 ---@param key string|string[] key for the mapping
 ---@param handler string | fun(): nil handler for the mapping
 ---@param opts? table<"'expr'"|"'noremap'"|"'nowait'"|"'remap'"|"'script'"|"'silent'"|"'unique'", boolean>
 ---@return nil
-function AvanteSplit:map(mode, key, handler, opts, ___force___) end
+function AvanteComp:map(mode, key, handler, opts, ___force___) end