feat(state): added support for restoring all window views and current…

… window

lua/flash/cache.lua

@@ -1,5 +1,6 @@
 local Pos = require("flash.search.pos")
 local Pattern = require("flash.search.pattern")
+local Util = require("flash.util")
 
 ---@class Flash.State.Window
 ---@field win number
@@ -41,6 +42,7 @@ function M:update()
   if self.state.win ~= win then
     self.state.win = win
     self.state.pos = Pos(vim.api.nvim_win_get_cursor(win))
+    self.state.restore_windows = Util.save_layout()
     M.cache = {}
     dirty = true
   end

lua/flash/state.lua

@@ -31,6 +31,7 @@ local Prompt = require("flash.prompt")
 ---@field visible boolean
 ---@field matcher fun(win: window, state:Flash.State): Flash.Matcher
 ---@field matchers Flash.Matcher[]
+---@field restore_windows? fun()
 ---@field ns number
 local M = {}
 M.__index = M
@@ -125,6 +126,13 @@ function M:jump(target)
   end
 end
 
+-- Will restore all window views
+function M:restore()
+  if self.restore_windows then
+    self.restore_windows()
+  end
+end
+
 function M:get_matcher(win)
   self.matchers[win] = self.matchers[win] or self.matcher(win, self)
   return self.matchers[win]

lua/flash/util.lua

@@ -16,4 +16,44 @@ function M.get_char()
   return ok and ret ~= M.ESC and ret or nil
 end
 
+function M.layout_wins()
+  local queue = { vim.fn.winlayout() }
+  ---@type table<window, window>
+  local wins = {}
+  while #queue > 0 do
+    local node = table.remove(queue)
+    if node[1] == "leaf" then
+      wins[node[2]] = node[2]
+    else
+      vim.list_extend(queue, node[2])
+    end
+  end
+  return wins
+end
+
+function M.save_layout()
+  local current_win = vim.api.nvim_get_current_win()
+  local wins = M.layout_wins()
+  ---@type table<window, table>
+  local state = {}
+  for _, win in pairs(wins) do
+    state[win] = vim.api.nvim_win_call(win, vim.fn.winsaveview)
+  end
+  return function()
+    for win, s in pairs(state) do
+      if vim.api.nvim_win_is_valid(win) then
+        local buf = vim.api.nvim_win_get_buf(win)
+        -- never restore terminal buffers to prevent flickering
+        if vim.bo[buf].buftype ~= "terminal" then
+          pcall(vim.api.nvim_win_call, win, function()
+            vim.fn.winrestview(s)
+          end)
+        end
+      end
+    end
+    vim.api.nvim_set_current_win(current_win)
+    state = {}
+  end
+end
+
 return M