From 6ef8c54fb526bf3a0bc4efb0b2fe8e6d9a7daed2 Mon Sep 17 00:00:00 2001
From: Lewis Russell <lewis6991@gmail.com>
Date: Wed, 29 Nov 2023 12:07:13 +0000
Subject: [PATCH] fix(preview): handle staged hunks

---
 lua/gitsigns/actions.lua | 51 +++++++++++++++++++++++++++-------------
 lua/gitsigns/debug.lua   |  1 -
 lua/gitsigns/manager.lua |  6 +++--
 3 files changed, 39 insertions(+), 19 deletions(-)

diff --git a/lua/gitsigns/actions.lua b/lua/gitsigns/actions.lua
index 997613b97..a894055a1 100644
--- a/lua/gitsigns/actions.lua
+++ b/lua/gitsigns/actions.lua
@@ -149,9 +149,10 @@ M.toggle_deleted = function(value)
   return config.show_deleted
 end
 
----@param bufnr? integer
----@param hunks? Gitsigns.Hunk.Hunk[]?
----@return Gitsigns.Hunk.Hunk?
+--- @param bufnr? integer
+--- @param hunks? Gitsigns.Hunk.Hunk[]?
+--- @return Gitsigns.Hunk.Hunk? hunk
+--- @return integer? index
 local function get_cursor_hunk(bufnr, hunks)
   bufnr = bufnr or current_buf()
 
@@ -185,13 +186,14 @@ local function get_range(params)
   return range
 end
 
+--- @async
 --- @param bufnr integer
 --- @param bcache Gitsigns.CacheEntry
 --- @param greedy? boolean
 --- @param staged? boolean
 --- @return Gitsigns.Hunk.Hunk[]? hunks
 local function get_hunks(bufnr, bcache, greedy, staged)
-  if greedy then
+  if greedy and config.diff_opts.linematch then
     -- Re-run the diff without linematch
     local buftext = util.buf_lines(bufnr)
     local text --- @type string[]?
@@ -225,7 +227,8 @@ local function get_hunk(bufnr, range, greedy, staged)
   local hunks = get_hunks(bufnr, bcache, greedy, staged)
 
   if not range then
-    return get_cursor_hunk(bufnr, hunks)
+    local hunk = get_cursor_hunk(bufnr, hunks)
+    return hunk
   end
 
   table.sort(range)
@@ -298,7 +301,7 @@ end
 --- @param bufnr integer
 --- @param hunk Gitsigns.Hunk.Hunk
 local function reset_hunk(bufnr, hunk)
-  local lstart, lend ---@type integer, integer
+  local lstart, lend --- @type integer, integer
   if hunk.type == 'delete' then
     lstart = hunk.added.start
     lend = hunk.added.start
@@ -732,11 +735,27 @@ local function feedkeys(keys)
   api.nvim_feedkeys(cy, 'n', false)
 end
 
+--- @param bufnr integer
+--- @param greedy? boolean
+--- @return Gitsigns.Hunk.Hunk? hunk
+--- @return boolean? staged
+local function get_hunk_with_staged(bufnr, greedy)
+  local hunk = get_hunk(bufnr, nil, greedy, false)
+  if hunk then
+    return hunk, false
+  end
+
+  hunk = get_hunk(bufnr, nil, greedy, true)
+  if hunk then
+    return hunk, true
+  end
+end
+
 --- Preview the hunk at the cursor position inline in the buffer.
-M.preview_hunk_inline = function()
+M.preview_hunk_inline = async.void(function()
   local bufnr = current_buf()
 
-  local hunk = get_cursor_hunk(bufnr)
+  local hunk, staged = get_hunk_with_staged(bufnr, true)
 
   if not hunk then
     return
@@ -744,11 +763,11 @@ M.preview_hunk_inline = function()
 
   clear_preview_inline(bufnr)
 
-  local winid ---@type integer
+  local winid --- @type integer
   manager.show_added(bufnr, ns_inline, hunk)
   if config._inline2 then
     if hunk.removed.count > 0 then
-      winid = manager.show_deleted_in_float(bufnr, ns_inline, hunk)
+      winid = manager.show_deleted_in_float(bufnr, ns_inline, hunk, staged)
     end
   else
     manager.show_deleted(bufnr, ns_inline, hunk)
@@ -771,7 +790,7 @@ M.preview_hunk_inline = function()
   if hunk.added.start <= 1 then
     feedkeys(hunk.removed.count .. '<C-y>')
   end
-end
+end)
 
 --- Select the hunk under the cursor.
 M.select_hunk = function()
@@ -929,8 +948,8 @@ C.blame_line = function(args, _)
   M.blame_line(args)
 end
 
----@param bcache Gitsigns.CacheEntry
----@param base string?
+--- @param bcache Gitsigns.CacheEntry
+--- @param base string?
 local function update_buf_base(bcache, base)
   bcache.base = base
   bcache:invalidate(true)
@@ -1121,14 +1140,14 @@ local function hunks_to_qflist(buf_or_filename, hunks, qflist)
   end
 end
 
----@param target 'all'|'attached'|integer|nil
----@return table[]?
+--- @param target 'all'|'attached'|integer|nil
+--- @return table[]?
 local function buildqflist(target)
   target = target or current_buf()
   if target == 0 then
     target = current_buf()
   end
-  local qflist = {} ---@type table[]
+  local qflist = {} --- @type table[]
 
   if type(target) == 'number' then
     local bufnr = target
diff --git a/lua/gitsigns/debug.lua b/lua/gitsigns/debug.lua
index a54105729..a6d659806 100644
--- a/lua/gitsigns/debug.lua
+++ b/lua/gitsigns/debug.lua
@@ -30,7 +30,6 @@ function M.dump_cache()
   --- @type string
   local text = vim.inspect(cache, { process = process })
   vim.api.nvim_echo({ { text } }, false, {})
-  return cache
 end
 
 --- @param noecho boolean
diff --git a/lua/gitsigns/manager.lua b/lua/gitsigns/manager.lua
index 3adc5a1cb..12912b667 100644
--- a/lua/gitsigns/manager.lua
+++ b/lua/gitsigns/manager.lua
@@ -294,8 +294,9 @@ end
 --- @param bufnr integer
 --- @param nsd integer
 --- @param hunk Gitsigns.Hunk.Hunk
+--- @param staged boolean?
 --- @return integer winid
-function M.show_deleted_in_float(bufnr, nsd, hunk)
+function M.show_deleted_in_float(bufnr, nsd, hunk, staged)
   local cwin = api.nvim_get_current_win()
   local virt_lines = {} --- @type {[1]: string, [2]: string}[][]
   local textoff = vim.fn.getwininfo(cwin)[1].textoff --[[@as integer]]
@@ -317,7 +318,8 @@ function M.show_deleted_in_float(bufnr, nsd, hunk)
 
   local bcache = cache[bufnr]
   local pbufnr = api.nvim_create_buf(false, true)
-  api.nvim_buf_set_lines(pbufnr, 0, -1, false, bcache.compare_text)
+  local text = staged and bcache.compare_text_head or bcache.compare_text
+  api.nvim_buf_set_lines(pbufnr, 0, -1, false, assert(text))
 
   local width = api.nvim_win_get_width(0)