refactor(git): rename stage_hunk to toggle_stage

lua/core/keymap.lua

@@ -235,7 +235,7 @@ vim.keymap.set("n", "<leader>gc", "<Plug>(git-commit)")
 vim.keymap.set("n", "<leader>ga", "<Plug>(git-commit-amend)")
 vim.keymap.set("n", "<leader>gl", "<Plug>(git-log)")
 vim.keymap.set("n", "<leader>gv", "<Plug>(git-hunk-select)")
-vim.keymap.set("n", "<leader>gs", "<Plug>(git-hunk-stage)")
+vim.keymap.set("n", "<leader>gs", "<Plug>(git-hunk-stage-toggle)")
 vim.keymap.set("n", "<leader>gr", "<Plug>(git-hunk-reset)")
 vim.keymap.set("n", "<C-w>g", "<Plug>(git-hunk-preview)")
 vim.keymap.set("n", "<leader>go", "<Plug>(git-overlay-toggle)")

lua/git/core/repo.lua

@@ -143,6 +143,7 @@ function Repo:_invalidate(relpath)
     end
     if affects_resolve(relpath) then
         self:_clear_cache_prefix("resolve:")
+        self:_clear_cache_prefix("head_blob:")
     end
     if relpath == "index" then
         self:_clear_cache_prefix("index:")
@@ -605,6 +606,15 @@ function Repo:index_sha(rel)
     return sha or nil
 end
 
+---@param rel string  worktree-relative path
+---@return string?
+function Repo:head_sha(rel)
+    local sha = self:get_cached("head_blob:" .. rel, function(self)
+        return self:rev_parse("HEAD:" .. rel, false) or false
+    end)
+    return sha or nil
+end
+
 ---@alias ow.Git.Repo.ResolveStatus "ok"|"ambiguous"|"missing"
 
 ---@param abbrev string

lua/git/hunks.lua

@@ -9,9 +9,9 @@ local NS_OVERLAY = vim.api.nvim_create_namespace("ow.git.hunks.overlay")
 ---@alias ow.Git.Hunks.HunkType "add"|"change"|"delete"
 
 ---@class ow.Git.Hunks.Hunk
----@field old_start integer
+---@field old_start integer  1-indexed first old line
 ---@field old_count integer
----@field new_start integer
+---@field new_start integer  1-indexed first new line
 ---@field new_count integer
 ---@field type ow.Git.Hunks.HunkType
 ---@field old_lines string[]
@@ -22,8 +22,11 @@ local NS_OVERLAY = vim.api.nvim_create_namespace("ow.git.hunks.overlay")
 ---@field rel string
 ---@field index string[]?
 ---@field index_sha string?
+---@field head string[]?
+---@field head_sha string?
 ---@field index_hl { src: string[], lines: table[][]? }?
 ---@field hunks ow.Git.Hunks.Hunk[]
+---@field staged ow.Git.Hunks.Hunk[]
 ---@field overlay boolean
 ---@field autocmds integer[]
 
@@ -64,17 +67,19 @@ local function diff_opts()
     return opts
 end
 
----@param state ow.Git.Hunks.BufState
+---@param old_lines string[]
 ---@param new_lines string[]
-local function compute_hunks(state, new_lines)
-    local old = table.concat(state.index or {}, "\n")
-    local new = table.concat(new_lines, "\n")
-    local raw = vim.text.diff(old, new, diff_opts())
+---@return ow.Git.Hunks.Hunk[]
+local function compute_hunks(old_lines, new_lines)
+    local raw = vim.text.diff(
+        table.concat(old_lines, "\n"),
+        table.concat(new_lines, "\n"),
+        diff_opts()
+    )
     ---@type ow.Git.Hunks.Hunk[]
     local hunks = {}
     if type(raw) ~= "table" then
-        state.hunks = hunks
-        return
+        return hunks
     end
     for _, h in ipairs(raw) do
         local os_ = h[1] --[[@as integer]]
@@ -89,16 +94,16 @@ local function compute_hunks(state, new_lines)
         else
             typ = "change"
         end
-        local old_lines = {}
-        if typ ~= "add" and state.index then
+        local old = {}
+        if typ ~= "add" then
             for i = os_, os_ + oc - 1 do
-                table.insert(old_lines, state.index[i] or "")
+                table.insert(old, old_lines[i] or "")
             end
         end
-        local hunk_new = {}
+        local new = {}
         if typ ~= "delete" then
             for i = ns_, ns_ + nc - 1 do
-                table.insert(hunk_new, new_lines[i] or "")
+                table.insert(new, new_lines[i] or "")
             end
         end
         table.insert(hunks, {
@@ -107,11 +112,11 @@ local function compute_hunks(state, new_lines)
             new_start = ns_,
             new_count = nc,
             type = typ,
-            old_lines = old_lines,
-            new_lines = hunk_new,
+            old_lines = old,
+            new_lines = new,
         })
     end
-    state.hunks = hunks
+    return hunks
 end
 
 ---@type table<ow.Git.Hunks.HunkType, string>
@@ -126,6 +131,95 @@ local function resolve_signs()
     return vim.tbl_extend("force", DEFAULT_SIGNS, cfg)
 end
 
+---@type table<ow.Git.Hunks.HunkType, string>
+local SIGN_HL = {
+    add = "GitHunkAdded",
+    change = "GitHunkChanged",
+    delete = "GitHunkRemoved",
+}
+
+---@type table<ow.Git.Hunks.HunkType, string>
+local STAGED_SIGN_HL = {
+    add = "GitHunkStagedAdded",
+    change = "GitHunkStagedChanged",
+    delete = "GitHunkStagedRemoved",
+}
+
+---@param h ow.Git.Hunks.Hunk
+---@param line_count integer
+---@return integer[]  0-indexed buffer rows for the hunk
+local function hunk_rows(h, line_count)
+    if h.type == "delete" then
+        local row = math.max(h.new_start, 1) - 1
+        if row >= line_count then
+            row = math.max(line_count - 1, 0)
+        end
+        return { row }
+    end
+    local rows = {}
+    for r = h.new_start, h.new_start + h.new_count - 1 do
+        local row = r - 1
+        if row >= 0 and row < line_count then
+            table.insert(rows, row)
+        end
+    end
+    return rows
+end
+
+---@param h ow.Git.Hunks.Hunk
+---@return integer  1-indexed last index line the hunk occupies
+local function index_end(h)
+    if h.old_count == 0 then
+        return h.old_start
+    end
+    return h.old_start + h.old_count - 1
+end
+
+---@param unstaged ow.Git.Hunks.Hunk[]
+---@param iline integer  1-indexed index line
+---@return integer?  1-indexed buffer line
+local function index_to_buffer(unstaged, iline)
+    local delta = 0
+    for _, h in ipairs(unstaged) do
+        if
+            h.old_count > 0
+            and iline >= h.old_start
+            and iline <= index_end(h)
+        then
+            return nil
+        end
+        if iline > index_end(h) then
+            delta = delta + h.new_count - h.old_count
+        end
+    end
+    return iline + delta
+end
+
+---@param state ow.Git.Hunks.BufState
+---@param line_count integer
+---@return { row: integer, hunk: ow.Git.Hunks.Hunk }[]  row is a 0-indexed buffer row
+local function staged_signs(state, line_count)
+    local out = {}
+    for _, h in ipairs(state.staged) do
+        local index_lines = {}
+        if h.type == "delete" then
+            table.insert(index_lines, math.max(h.new_start, 1))
+        else
+            for i = h.new_start, h.new_start + h.new_count - 1 do
+                table.insert(index_lines, i)
+            end
+        end
+        for _, iline in ipairs(index_lines) do
+            local bline = index_to_buffer(state.hunks, iline)
+            if bline then
+                local row = math.min(math.max(bline - 1, 0), line_count - 1)
+                table.insert(out, { row = row, hunk = h })
+            end
+        end
+    end
+    return out
+end
+
 ---@param buf integer
 local function render_signs(buf)
     if not vim.api.nvim_buf_is_valid(buf) then
@@ -138,37 +232,24 @@ local function render_signs(buf)
     end
     local signs = resolve_signs()
     local line_count = vim.api.nvim_buf_line_count(buf)
+    local signed = {}
     for _, h in ipairs(state.hunks) do
-        if h.type == "delete" then
-            local row = math.max(h.new_start, 1) - 1
-            if row >= line_count then
-                row = math.max(line_count - 1, 0)
-            end
+        for _, row in ipairs(hunk_rows(h, line_count)) do
+            signed[row] = true
             pcall(vim.api.nvim_buf_set_extmark, buf, NS_SIGNS, row, 0, {
-                sign_text = signs.delete,
-                sign_hl_group = "GitHunkRemoved",
+                sign_text = signs[h.type],
+                sign_hl_group = SIGN_HL[h.type],
                 priority = 100,
             })
-        else
-            local hl = h.type == "add" and "GitHunkAdded" or "GitHunkChanged"
-            local sign = h.type == "add" and signs.add or signs.change
-            for r = h.new_start, h.new_start + h.new_count - 1 do
-                local row = r - 1
-                if row >= 0 and row < line_count then
-                    pcall(
-                        vim.api.nvim_buf_set_extmark,
-                        buf,
-                        NS_SIGNS,
-                        row,
-                        0,
-                        {
-                            sign_text = sign,
-                            sign_hl_group = hl,
+        end
+    end
+    for _, s in ipairs(staged_signs(state, line_count)) do
+        if not signed[s.row] then
+            pcall(vim.api.nvim_buf_set_extmark, buf, NS_SIGNS, s.row, 0, {
+                sign_text = signs[s.hunk.type],
+                sign_hl_group = STAGED_SIGN_HL[s.hunk.type],
                 priority = 100,
-                        }
-                    )
-                end
-            end
+            })
         end
     end
 end
@@ -346,27 +427,32 @@ end
 
 ---@param state ow.Git.Hunks.BufState
 ---@param buf integer
----@param new_sha string
-local function load_index_and_render(state, buf, new_sha)
-    util.git({ "cat-file", "-p", ":0:" .. state.rel }, {
+---@param rev string
+---@param want string?  the wanted blob sha
+---@param have string?  the currently-loaded blob sha
+---@param apply fun(lines: string[]?, sha: string?)
+---@param after fun()
+local function ensure_content(state, buf, rev, want, have, apply, after)
+    if not want then
+        apply(nil, nil)
+        return after()
+    end
+    if want == have then
+        return after()
+    end
+    util.git({ "cat-file", "-p", rev }, {
         cwd = state.repo.worktree,
         silent = true,
         on_exit = function(res)
             if states[buf] ~= state or not vim.api.nvim_buf_is_valid(buf) then
                 return
             end
-            if res.code ~= 0 then
-                state.index = {}
-                state.index_sha = nil
-                state.hunks = {}
-                render(buf)
-                return
+            if res.code == 0 then
+                apply(util.split_lines(res.stdout or ""), want)
+            else
+                apply(nil, nil)
             end
-            state.index = util.split_lines(res.stdout or "")
-            state.index_sha = new_sha
-            local new_lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
-            compute_hunks(state, new_lines)
-            render(buf)
+            after()
         end,
     })
 end
@@ -380,21 +466,43 @@ local function recompute(buf)
     if not state then
         return
     end
-    local new_sha = state.repo:index_sha(state.rel)
-    if not new_sha then
-        state.index = nil
-        state.index_sha = nil
-        state.hunks = {}
+    local r = state.repo
+    ensure_content(
+        state,
+        buf,
+        ":0:" .. state.rel,
+        r:index_sha(state.rel),
+        state.index_sha,
+        function(lines, sha)
+            state.index = lines
+            state.index_sha = sha
+        end,
+        function()
+            ensure_content(
+                state,
+                buf,
+                "HEAD:" .. state.rel,
+                r:head_sha(state.rel),
+                state.head_sha,
+                function(lines, sha)
+                    state.head = lines
+                    state.head_sha = sha
+                end,
+                function()
+                    local new =
+                        vim.api.nvim_buf_get_lines(buf, 0, -1, false)
+                    state.hunks = state.index
+                            and compute_hunks(state.index, new)
+                        or {}
+                    state.staged = state.head
+                            and state.index
+                            and compute_hunks(state.head, state.index)
+                        or {}
                     render(buf)
-        return
                 end
-    if new_sha == state.index_sha and state.index then
-        local new_lines = vim.api.nvim_buf_get_lines(buf, 0, -1, false)
-        compute_hunks(state, new_lines)
-        render(buf)
-        return
+            )
         end
-    load_index_and_render(state, buf, new_sha)
+    )
 end
 
 local schedule, sched_handle = util.keyed_debounce(recompute, 100)
@@ -426,7 +534,10 @@ function M.attach(buf)
         rel = rel,
         index = nil,
         index_sha = nil,
+        head = nil,
+        head_sha = nil,
         hunks = {},
+        staged = {},
         overlay = vim.g.git_hunk_overlay_default == true,
         autocmds = {},
     }
@@ -488,28 +599,32 @@ function M.toggle_overlay(buf)
     render(buf)
 end
 
----@param buf? integer
+---@param hunks ow.Git.Hunks.Hunk[]
 ---@param row integer 1-indexed cursor line
 ---@return ow.Git.Hunks.Hunk?
-local function hunk_at(buf, row)
-    buf = buf or vim.api.nvim_get_current_buf()
-    local state = states[buf]
-    if not state then
-        return nil
-    end
-    for _, h in ipairs(state.hunks) do
+local function hunk_at(hunks, row)
+    for _, h in ipairs(hunks) do
         if h.type == "delete" then
-            local anchor = math.max(h.new_start, 1)
-            if anchor == row then
+            if math.max(h.new_start, 1) == row then
                 return h
             end
-        else
-            local lo = h.new_start
-            local hi = h.new_start + h.new_count - 1
-            if row >= lo and row <= hi then
+        elseif row >= h.new_start and row <= h.new_start + h.new_count - 1 then
             return h
         end
     end
+    return nil
+end
+
+---@param state ow.Git.Hunks.BufState
+---@param buf integer
+---@param row integer 1-indexed cursor line
+---@return ow.Git.Hunks.Hunk?
+local function staged_hunk_at(state, buf, row)
+    local line_count = vim.api.nvim_buf_line_count(buf)
+    for _, s in ipairs(staged_signs(state, line_count)) do
+        if s.row == row - 1 then
+            return s.hunk
+        end
     end
     return nil
 end
@@ -524,11 +639,11 @@ local function cursor_hunk(buf)
     if not state then
         return buf, nil, nil
     end
-    return buf, state, hunk_at(buf, vim.api.nvim_win_get_cursor(0)[1])
+    return buf, state, hunk_at(state.hunks, vim.api.nvim_win_get_cursor(0)[1])
 end
 
 ---@param h ow.Git.Hunks.Hunk
----@return integer
+---@return integer  1-indexed buffer line to anchor the cursor on
 local function anchor_line(h)
     if h.type == "delete" then
         return math.max(h.new_start, 1)
@@ -588,39 +703,138 @@ local function hunk_body(h)
     return lines
 end
 
+local PATCH_CONTEXT = 3
+
 ---@param h ow.Git.Hunks.Hunk
----@param state ow.Git.Hunks.BufState
----@return string
-local function build_patch(h, state)
-    local lines = { "--- a/" .. state.rel, "+++ b/" .. state.rel }
-    vim.list_extend(lines, hunk_body(h))
-    return table.concat(lines, "\n") .. "\n"
+---@return integer old_before  count of old lines before the hunk's changed content
+---@return integer new_before  count of new lines before the hunk's changed content
+local function hunk_offsets(h)
+    if h.type == "add" then
+        return h.old_start, h.new_start - 1
+    elseif h.type == "delete" then
+        return h.old_start - 1, h.new_start
+    end
+    return h.old_start - 1, h.new_start - 1
 end
 
----@param buf? integer
-function M.stage_hunk(buf)
-    local _, state, h = cursor_hunk(buf)
-    if not state then
-        return
+---@param h ow.Git.Hunks.Hunk
+---@return ow.Git.Hunks.Hunk
+local function invert(h)
+    local typ ---@type ow.Git.Hunks.HunkType
+    if h.type == "add" then
+        typ = "delete"
+    elseif h.type == "delete" then
+        typ = "add"
+    else
+        typ = "change"
     end
-    if not h then
-        util.warning("git hunks: no hunk at cursor")
-        return
+    return {
+        old_start = h.new_start,
+        old_count = h.new_count,
+        new_start = h.old_start,
+        new_count = h.old_count,
+        type = typ,
+        old_lines = h.new_lines,
+        new_lines = h.old_lines,
+    }
+end
+
+---@param h ow.Git.Hunks.Hunk
+---@param old_lines string[]
+---@param rel string
+---@return string patch
+---@return boolean zero_context
+local function build_patch(h, old_lines, rel)
+    local old_before, new_before = hunk_offsets(h)
+    local pre = {}
+    for i = math.max(old_before - PATCH_CONTEXT + 1, 1), old_before do
+        pre[#pre + 1] = old_lines[i] or ""
     end
-    util.git({ "apply", "--cached", "--unidiff-zero", "-" }, {
+    local post = {}
+    local after = old_before + h.old_count
+    for i = after + 1, math.min(after + PATCH_CONTEXT, #old_lines) do
+        post[#post + 1] = old_lines[i] or ""
+    end
+    local old_n = #pre + h.old_count + #post
+    local new_n = #pre + h.new_count + #post
+    local old_start = old_n > 0 and old_before - #pre + 1 or old_before
+    local new_start = new_n > 0 and new_before - #pre + 1 or new_before
+    local body = {
+        string.format(
+            "@@ -%d,%d +%d,%d @@",
+            old_start,
+            old_n,
+            new_start,
+            new_n
+        ),
+    }
+    for _, l in ipairs(pre) do
+        body[#body + 1] = " " .. l
+    end
+    for _, l in ipairs(h.old_lines) do
+        body[#body + 1] = "-" .. l
+    end
+    for _, l in ipairs(h.new_lines) do
+        body[#body + 1] = "+" .. l
+    end
+    for _, l in ipairs(post) do
+        body[#body + 1] = " " .. l
+    end
+    local lines = { "--- a/" .. rel, "+++ b/" .. rel }
+    vim.list_extend(lines, body)
+    return table.concat(lines, "\n") .. "\n", #pre == 0 and #post == 0
+end
+
+---@param state ow.Git.Hunks.BufState
+---@param buf integer
+---@param patch string
+---@param zero_context boolean
+local function apply_patch(state, buf, patch, zero_context)
+    local args = { "apply", "--cached" }
+    if zero_context then
+        table.insert(args, "--unidiff-zero")
+    end
+    table.insert(args, "-")
+    util.git(args, {
         cwd = state.repo.worktree,
-        stdin = build_patch(h, state),
+        stdin = patch,
         on_exit = function(res)
             if res.code ~= 0 then
-                util.error(
-                    "git apply failed: %s",
-                    vim.trim(res.stderr or "")
-                )
+                util.error("git apply failed: %s", vim.trim(res.stderr or ""))
+                return
+            end
+            local s = states[buf]
+            if s then
+                s.index_sha = nil
+                schedule(buf)
             end
         end,
     })
 end
 
+---@param buf? integer
+function M.toggle_stage(buf)
+    buf = resolve_buf(buf)
+    local state = states[buf]
+    if not state then
+        return
+    end
+    local row = vim.api.nvim_win_get_cursor(0)[1]
+    local unstaged = hunk_at(state.hunks, row)
+    if unstaged and state.index then
+        local patch, zero = build_patch(unstaged, state.index, state.rel)
+        apply_patch(state, buf, patch, zero)
+        return
+    end
+    local staged = staged_hunk_at(state, buf, row)
+    if staged and state.index then
+        local patch, zero = build_patch(invert(staged), state.index, state.rel)
+        apply_patch(state, buf, patch, zero)
+        return
+    end
+    util.warning("git hunks: no hunk at cursor")
+end
+
 ---@param buf? integer
 function M.reset_hunk(buf)
     local target, state, h = cursor_hunk(buf)

plugin/git.lua

@@ -41,12 +41,47 @@ local DEFAULT_HIGHLIGHTS = {
     GitHunkAddLine = "DiffAdd",
     GitHunkDeleteLine = "DiffDelete",
 }
-for name, link in pairs(DEFAULT_HIGHLIGHTS) do
-    vim.api.nvim_set_hl(0, name, { link = link, default = true })
+local STAGED_HUNK_HL = {
+    GitHunkStagedAdded = "GitHunkAdded",
+    GitHunkStagedChanged = "GitHunkChanged",
+    GitHunkStagedRemoved = "GitHunkRemoved",
+}
+
+local function blend(a, b, t)
+    local function mix(shift)
+        local x = bit.band(bit.rshift(a, shift), 0xff)
+        local y = bit.band(bit.rshift(b, shift), 0xff)
+        return bit.lshift(math.floor(x + (y - x) * t + 0.5), shift)
+    end
+    return mix(16) + mix(8) + mix(0)
 end
 
+local function apply_highlights()
+    for name, link in pairs(DEFAULT_HIGHLIGHTS) do
+        vim.api.nvim_set_hl(0, name, { link = link, default = true })
+    end
+    local bg = vim.api.nvim_get_hl(0, { name = "Normal" }).bg or 0x000000
+    for name, base in pairs(STAGED_HUNK_HL) do
+        local src = vim.api.nvim_get_hl(0, { name = base, link = false })
+        local hl = {}
+        if src.fg then
+            hl.fg = blend(src.fg, bg, 0.45)
+        end
+        if src.bg then
+            hl.bg = blend(src.bg, bg, 0.45)
+        end
+        vim.api.nvim_set_hl(0, name, hl)
+    end
+end
+apply_highlights()
+
 local group = vim.api.nvim_create_augroup("ow.git", { clear = true })
 
+vim.api.nvim_create_autocmd("ColorScheme", {
+    group = group,
+    callback = apply_highlights,
+})
+
 vim.api.nvim_create_autocmd({ "BufReadPost", "BufNewFile" }, {
     group = group,
     callback = function(args)
@@ -260,9 +295,9 @@ end, { silent = true, desc = "Jump to next git hunk" })
 vim.keymap.set({ "n", "x" }, "<Plug>(git-hunk-prev)", function()
     require("git.hunks").nav("prev")
 end, { silent = true, desc = "Jump to previous git hunk" })
-vim.keymap.set("n", "<Plug>(git-hunk-stage)", function()
-    require("git.hunks").stage_hunk()
-end, { silent = true, desc = "Stage hunk under cursor" })
+vim.keymap.set("n", "<Plug>(git-hunk-stage-toggle)", function()
+    require("git.hunks").toggle_stage()
+end, { silent = true, desc = "Stage or unstage the hunk under cursor" })
 vim.keymap.set("n", "<Plug>(git-hunk-reset)", function()
     require("git.hunks").reset_hunk()
 end, { silent = true, desc = "Reset hunk under cursor" })

plugins/onedark.lua

@@ -45,10 +45,8 @@ local highlights = {
     TabLineFill = { bg = c.bg1 },
     EndOfBuffer = { fg = "NONE", bg = "NONE" },
     DiffAdd = { bg = "#1a2f22" },
-    DiffChange = { bg = "#15304a" },
-    DiffDelete = { bg = "#311c1e" },
-    Changed = { fg = c.yellow },
     NvimTreeIndentMarker = { fg = c.bg3 },
+    GitUnstaged = { fg = c.yellow },
 }
 for kind, color in pairs(completion_kind_colors) do
     highlights["LspKind" .. kind] = { fg = color }

test/git/hunks_test.lua

@@ -18,8 +18,8 @@ local function setup(committed, worktree, file)
     hunks._flush(buf)
     t.wait_for(function()
         local s = hunks.state(buf)
-        return s ~= nil and s.index ~= nil
-    end, "hunks to compute the index snapshot")
+        return s ~= nil and s.index ~= nil and s.head ~= nil
+    end, "hunks to load the index and HEAD snapshots")
     local state = assert(hunks.state(buf), "buffer state should exist")
     return dir, buf, state
 end
@@ -261,10 +261,10 @@ t.test("overlay: toggling swaps gutter signs for the overlay", function()
     )
 end)
 
-t.test("stage_hunk stages the change into the index", function()
+t.test("toggle_stage stages the change into the index", function()
     local dir, buf = setup("a\nb\nc\n", "a\nB\nc\n")
     vim.api.nvim_win_set_cursor(0, { 2, 0 })
-    hunks.stage_hunk(buf)
+    hunks.toggle_stage(buf)
     t.wait_for(function()
         return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
     end, "stage to land in the index")
@@ -276,26 +276,256 @@ t.test("stage_hunk stages the change into the index", function()
     )
 end)
 
-t.test("stage_hunk stages a pure addition", function()
+t.test("toggle_stage stages a pure addition", function()
     local dir, buf = setup("a\nb\n", "a\nb\nc\n")
     vim.api.nvim_win_set_cursor(0, { 3, 0 })
-    hunks.stage_hunk(buf)
+    hunks.toggle_stage(buf)
     t.wait_for(function()
         return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
     end, "stage to land in the index")
     t.eq(h.git(dir, "show", ":0:a.txt").stdout, "a\nb\nc")
 end)
 
-t.test("stage_hunk stages a deletion", function()
+t.test("toggle_stage stages a deletion", function()
     local dir, buf = setup("a\nb\nc\n", "a\nc\n")
     vim.api.nvim_win_set_cursor(0, { 1, 0 })
-    hunks.stage_hunk(buf)
+    hunks.toggle_stage(buf)
     t.wait_for(function()
         return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
     end, "stage to land in the index")
     t.eq(h.git(dir, "show", ":0:a.txt").stdout, "a\nc")
 end)
 
+t.test("toggle_stage stages only the hunk under the cursor", function()
+    local committed = table.concat({
+        "local M = {}",
+        "",
+        "function M.first()",
+        "    return 1",
+        "end",
+        "",
+        "function M.last()",
+        "    return 9",
+        "end",
+        "",
+        "return M",
+    }, "\n") .. "\n"
+    local worktree = table.concat({
+        "local M = {}",
+        "",
+        "-- helpers",
+        "function M.first()",
+        "    return 1",
+        "end",
+        "",
+        "function M.mid()",
+        "    return 5",
+        "end",
+        "",
+        "function M.last()",
+        "    return 9",
+        "end",
+        "",
+        "return M",
+    }, "\n") .. "\n"
+    local dir, buf = setup(committed, worktree)
+    vim.api.nvim_win_set_cursor(0, { 9, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
+    end, "the mid hunk to land in the index")
+    t.eq(
+        h.git(dir, "show", ":0:a.txt").stdout,
+        table.concat({
+            "local M = {}",
+            "",
+            "function M.first()",
+            "    return 1",
+            "end",
+            "",
+            "function M.mid()",
+            "    return 5",
+            "end",
+            "",
+            "function M.last()",
+            "    return 9",
+            "end",
+            "",
+            "return M",
+        }, "\n"),
+        "only the cursor's hunk is staged, placed at the right line"
+    )
+end)
+
+t.test("toggle_stage stages a whole-file change with no context", function()
+    local dir, buf = setup("a\nb\nc\n", "x\ny\nz\n")
+    vim.api.nvim_win_set_cursor(0, { 2, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
+    end, "the change to land in the index")
+    t.eq(h.git(dir, "show", ":0:a.txt").stdout, "x\ny\nz")
+end)
+
+t.test("toggle_stage stages a change at the start of the file", function()
+    local dir, buf = setup("a\nb\nc\nd\ne\n", "A\nb\nc\nd\ne\n")
+    vim.api.nvim_win_set_cursor(0, { 1, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
+    end, "the change to land in the index")
+    t.eq(h.git(dir, "show", ":0:a.txt").stdout, "A\nb\nc\nd\ne")
+end)
+
+t.test("toggle_stage stages a change at the end of the file", function()
+    local dir, buf = setup("a\nb\nc\nd\ne\n", "a\nb\nc\nd\nE\n")
+    vim.api.nvim_win_set_cursor(0, { 5, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
+    end, "the change to land in the index")
+    t.eq(h.git(dir, "show", ":0:a.txt").stdout, "a\nb\nc\nd\nE")
+end)
+
+t.test("toggle_stage stages a deletion at the start of the file", function()
+    local dir, buf = setup("a\nb\nc\nd\n", "b\nc\nd\n")
+    vim.api.nvim_win_set_cursor(0, { 1, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
+    end, "the deletion to land in the index")
+    t.eq(h.git(dir, "show", ":0:a.txt").stdout, "b\nc\nd")
+end)
+
+t.test("toggle_stage leaves an adjacent unstaged hunk in place", function()
+    local dir, buf = setup("a\nb\nc\nd\ne\n", "A\nb\nC\nd\ne\n")
+    vim.api.nvim_win_set_cursor(0, { 3, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return h.git(dir, "diff", "--cached", "--name-only").stdout ~= ""
+    end, "the line-3 hunk to land in the index")
+    t.eq(
+        h.git(dir, "show", ":0:a.txt").stdout,
+        "a\nb\nC\nd\ne",
+        "only line 3 is staged; the adjacent line-1 hunk is untouched"
+    )
+end)
+
+t.test("toggle_stage unstages one of two adjacent staged hunks", function()
+    local dir, buf = setup("a\nb\nc\nd\ne\n", "A\nb\nC\nd\ne\n")
+    vim.api.nvim_win_set_cursor(0, { 1, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).staged == 1
+    end, "the line-1 hunk to be staged")
+    vim.api.nvim_win_set_cursor(0, { 3, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        local s = assert(hunks.state(buf))
+        return #s.hunks == 0 and #s.staged == 2
+    end, "both hunks to be staged")
+
+    vim.api.nvim_win_set_cursor(0, { 3, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).staged == 1
+    end, "the line-3 hunk to be unstaged again")
+    t.eq(
+        h.git(dir, "show", ":0:a.txt").stdout,
+        "A\nb\nc\nd\ne",
+        "line 3 reverts to HEAD while the staged line-1 change remains"
+    )
+end)
+
+t.test("toggle_stage refreshes the gutter when status stays modified", function()
+    local _, buf = setup("a\nb\nc\nd\ne\n", "A\nb\nC\nd\nE\n")
+    t.eq(#assert(hunks.state(buf)).hunks, 3)
+
+    vim.api.nvim_win_set_cursor(0, { 1, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).hunks == 2
+    end, "gutter to drop the first staged hunk")
+
+    vim.api.nvim_win_set_cursor(0, { 3, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).hunks == 1
+    end, "gutter to drop the middle staged hunk")
+end)
+
+t.test("staged hunks show with the staged highlight", function()
+    local _, buf = setup("a\nb\nc\n", "a\nB\nc\n")
+    vim.api.nvim_win_set_cursor(0, { 2, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        local s = assert(hunks.state(buf))
+        return #s.hunks == 0 and #s.staged == 1
+    end, "the hunk to move from unstaged to staged")
+    t.eq(sign_marks(buf), {
+        { row = 1, sign = "┃", hl = "GitHunkStagedChanged" },
+    })
+end)
+
+t.test("the gutter shows staged and unstaged hunks together", function()
+    local _, buf = setup("a\nb\nc\nd\ne\n", "A\nb\nC\nd\nE\n")
+    vim.api.nvim_win_set_cursor(0, { 1, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).hunks == 2
+    end, "the first hunk to leave the unstaged set")
+    t.eq(sign_marks(buf), {
+        { row = 0, sign = "┃", hl = "GitHunkStagedChanged" },
+        { row = 2, sign = "┃", hl = "GitHunkChanged" },
+        { row = 4, sign = "┃", hl = "GitHunkChanged" },
+    })
+end)
+
+t.test("toggle_stage toggles a staged hunk back to unstaged", function()
+    local _, buf = setup("a\nb\nc\n", "a\nB\nc\n")
+    vim.api.nvim_win_set_cursor(0, { 2, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        local s = assert(hunks.state(buf))
+        return #s.hunks == 0 and #s.staged == 1
+    end, "the hunk to become staged")
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        local s = assert(hunks.state(buf))
+        return #s.hunks == 1 and #s.staged == 0
+    end, "the hunk to return to unstaged")
+    t.eq(sign_marks(buf), {
+        { row = 1, sign = "┃", hl = "GitHunkChanged" },
+    })
+end)
+
+t.test("toggle_stage unstages correctly when buffer lines are shifted", function()
+    local dir, buf = setup("a\nb\nc\n", "a\nb\nC\n")
+    vim.api.nvim_win_set_cursor(0, { 3, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).staged == 1
+    end, "the line-3 change to be staged")
+
+    vim.api.nvim_buf_set_lines(buf, 0, 0, false, { "NEW" })
+    vim.api.nvim_exec_autocmds("TextChanged", { buffer = buf })
+    hunks._flush(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).hunks == 1
+    end, "the unstaged add at the top to register")
+
+    vim.api.nvim_win_set_cursor(0, { 4, 0 })
+    hunks.toggle_stage(buf)
+    t.wait_for(function()
+        return #assert(hunks.state(buf)).staged == 0
+    end, "the shifted staged hunk to be unstaged")
+    t.eq(
+        h.git(dir, "show", ":0:a.txt").stdout,
+        "a\nb\nc",
+        "the index reverts to HEAD content for the unstaged hunk"
+    )
+end)
+
 t.test("reset_hunk restores the index content for a change", function()
     local _, buf, state = setup("a\nb\nc\n", "a\nB\nc\n")
     vim.api.nvim_win_set_cursor(0, { 2, 0 })

test/git/repo_test.lua

@@ -97,6 +97,25 @@ t.test("index_sha cache clears when the index is written", function()
     wait_cleared(r, "index:a", 2000)
 end)
 
+t.test("head_sha returns the blob sha and caches it", function()
+    local dir = h.make_repo({ a = "x\n" })
+    local r = assert(require("git.core.repo").resolve(dir))
+    local sha = r:head_sha("a")
+    t.truthy(sha and #sha > 0, "head_sha returns the HEAD blob sha")
+    t.truthy(r._cache["head_blob:a"] ~= nil, "the result is cached")
+    t.eq(r:head_sha("a"), sha, "a cached call returns the same sha")
+end)
+
+t.test("head_sha cache clears when HEAD moves", function()
+    local dir = h.make_repo({ a = "x\n" })
+    local r = assert(require("git.core.repo").resolve(dir))
+    r:head_sha("a")
+    t.truthy(r._cache["head_blob:a"] ~= nil, "sha is cached before the commit")
+    t.write(dir, "a", "y\n")
+    h.git(dir, "commit", "-aqm", "change")
+    wait_cleared(r, "head_blob:a", 2000)
+end)
+
 t.test("cache clears after top-level .git change (commit)", function()
     local dir = h.make_repo({ a = "x" })
     local r = assert(require("git.core.repo").resolve(dir))