From fd331afe6408bd9c6781959dbf4e88bcaa8ebd02 Mon Sep 17 00:00:00 2001
From: Akrm Al-Hakimi <alk9fh@bosch.com>
Date: Mon, 20 Jul 2026 08:37:33 -0400
Subject: [PATCH] fix: restore hunk review state on undo and redo

Snapshot pending hunk state when accepting changes and restore matching
rendering when buffer contents return through undo or redo. Clear stale
snapshots when new proposals replace an existing review session.

Closes #3
---
 lua/jumpy/navigate.lua | 99 +++++++++++++++++++++++++++++++++++++++++-
 lua/jumpy/prompt.lua   |  2 +
 lua/jumpy/render.lua   | 55 +++++++++++++++++++++--
 3 files changed, 151 insertions(+), 5 deletions(-)

diff --git a/lua/jumpy/navigate.lua b/lua/jumpy/navigate.lua
index 2647953..f1192f0 100644
--- a/lua/jumpy/navigate.lua
+++ b/lua/jumpy/navigate.lua
@@ -4,6 +4,86 @@ local render = require("jumpy.render")
 
 local MSG_NO_HUNKS = "jumpy: no hunks"
 local MSG_NO_HUNK_UNDER_CURSOR = "jumpy: no hunk under cursor"
+local offset_table = {}
+local undo_history = {}
+local syncing_undo = {}
+
+local function copy_offsets(bufnr)
+  return vim.deepcopy(offset_table[bufnr] or {})
+end
+
+local function restore_offsets(bufnr, offsets)
+  offset_table[bufnr] = vim.deepcopy(offsets or {})
+end
+
+local function same_lines(a, b)
+  return vim.deep_equal(a, b)
+end
+
+local function sync_undo_state(bufnr)
+  if syncing_undo[bufnr] or not vim.api.nvim_buf_is_valid(bufnr) then
+    return
+  end
+
+  local history = undo_history[bufnr]
+  if not history then
+    return
+  end
+
+  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
+  for i = #history, 1, -1 do
+    local entry = history[i]
+    if same_lines(lines, entry.before_lines) then
+      syncing_undo[bufnr] = true
+      render.restore(bufnr, entry.before_state)
+      restore_offsets(bufnr, entry.before_offsets)
+      syncing_undo[bufnr] = nil
+      M._refresh_quickfix()
+      return
+    end
+    if same_lines(lines, entry.after_lines) then
+      syncing_undo[bufnr] = true
+      render.restore(bufnr, entry.after_state)
+      restore_offsets(bufnr, entry.after_offsets)
+      syncing_undo[bufnr] = nil
+      M._refresh_quickfix()
+      return
+    end
+  end
+end
+
+M._sync_undo_state = sync_undo_state
+
+local function record_undo_state(bufnr, before_state, before_offsets, before_lines)
+  local history = undo_history[bufnr] or {}
+  table.insert(history, {
+    before_state = before_state,
+    before_offsets = before_offsets,
+    before_lines = before_lines,
+    after_state = render.snapshot(bufnr),
+    after_offsets = copy_offsets(bufnr),
+    after_lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false),
+  })
+  undo_history[bufnr] = history
+end
+
+function M._clear_undo_history(bufnr)
+  undo_history[bufnr] = nil
+  offset_table[bufnr] = nil
+  syncing_undo[bufnr] = nil
+end
+
+vim.api.nvim_create_autocmd("TextChanged", {
+  callback = function(args)
+    sync_undo_state(args.buf)
+  end,
+})
+
+vim.api.nvim_create_autocmd("BufWipeout", {
+  callback = function(args)
+    M._clear_undo_history(args.buf)
+  end,
+})
 
 local function center_cursor()
   vim.cmd("normal! zz")
@@ -113,16 +193,25 @@ function M.accept_hunk(bufnr, hunk_idx)
   end
 
   local offset = M._get_offset(bufnr, hunk_idx)
+  local before_state = render.snapshot(bufnr)
+  local before_offsets = copy_offsets(bufnr)
+  local before_lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
 
   local start_line = hunk.old_start - 1 + offset
   local end_line = start_line + hunk.old_count
 
+  syncing_undo[bufnr] = true
   vim.api.nvim_buf_set_lines(bufnr, start_line, end_line, false, hunk.added_lines)
 
   local delta = #hunk.added_lines - hunk.old_count
   M._apply_offset(bufnr, hunk_idx, delta)
 
   render.clear_hunk(bufnr, hunk_idx)
+  if not render.get_state(bufnr) then
+    offset_table[bufnr] = nil
+  end
+  record_undo_state(bufnr, before_state, before_offsets, before_lines)
+  syncing_undo[bufnr] = nil
   M._refresh_quickfix()
   return true
 end
@@ -177,6 +266,11 @@ function M.accept_all()
     table.insert(reversed, active[i])
   end
 
+  local before_state = render.snapshot(bufnr)
+  local before_offsets = copy_offsets(bufnr)
+  local before_lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
+
+  syncing_undo[bufnr] = true
   for _, entry in ipairs(reversed) do
     local hunk = entry.hunk
     local start_line = hunk.old_start - 1
@@ -185,6 +279,9 @@ function M.accept_all()
   end
 
   render.clear(bufnr)
+  offset_table[bufnr] = nil
+  record_undo_state(bufnr, before_state, before_offsets, before_lines)
+  syncing_undo[bufnr] = nil
   M._refresh_quickfix()
   vim.notify("jumpy: all hunks accepted", vim.log.levels.INFO)
 end
@@ -336,8 +433,6 @@ function M._setup_quickfix_keymaps()
   })
 end
 
-local offset_table = {}
-
 function M._get_offset(bufnr, hunk_idx)
   if not offset_table[bufnr] then
     offset_table[bufnr] = {}
diff --git a/lua/jumpy/prompt.lua b/lua/jumpy/prompt.lua
index 4f9bad8..747ec7e 100644
--- a/lua/jumpy/prompt.lua
+++ b/lua/jumpy/prompt.lua
@@ -365,6 +365,7 @@ function M._submit()
               if bufnr then
                 local hunks = diff.compute(file.lines, result.lines)
                 if #hunks > 0 then
+                  require("jumpy.navigate")._clear_undo_history(bufnr)
                   render.show(bufnr, hunks, file.lines, result.lines)
                   total_hunks = total_hunks + #hunks
                 end
@@ -427,6 +428,7 @@ function M._submit()
             return
           end
 
+          require("jumpy.navigate")._clear_undo_history(source_buf)
           render.show(source_buf, hunks, source_lines, proposed_lines)
 
           vim.notify(string.format("jumpy: %d hunk(s) proposed", #hunks), vim.log.levels.INFO)
diff --git a/lua/jumpy/render.lua b/lua/jumpy/render.lua
index 6bf7964..47111a6 100644
--- a/lua/jumpy/render.lua
+++ b/lua/jumpy/render.lua
@@ -19,6 +19,21 @@ end
 
 local buf_states = {}
 
+local function copy_lines(lines)
+  return vim.deepcopy(lines)
+end
+
+local function copy_hunk(hunk)
+  return {
+    removed_lines = copy_lines(hunk.removed_lines),
+    added_lines = copy_lines(hunk.added_lines),
+    old_start = hunk.old_start,
+    old_count = hunk.old_count,
+    new_start = hunk.new_start,
+    new_count = hunk.new_count,
+  }
+end
+
 function M.get_state(bufnr)
   return buf_states[bufnr]
 end
@@ -37,6 +52,26 @@ function M.get_all_states()
   return states
 end
 
+function M.snapshot(bufnr)
+  local state = buf_states[bufnr]
+  if not state then
+    return nil
+  end
+
+  local hunks = {}
+  for idx, hunk in pairs(state.hunks) do
+    if hunk then
+      hunks[idx] = copy_hunk(hunk)
+    end
+  end
+
+  return {
+    hunks = hunks,
+    original_lines = copy_lines(state.original_lines),
+    proposed_lines = copy_lines(state.proposed_lines),
+  }
+end
+
 function M.show(bufnr, hunks, original_lines, proposed_lines)
   M.clear(bufnr)
 
@@ -47,9 +82,16 @@ function M.show(bufnr, hunks, original_lines, proposed_lines)
     extmark_ids = {},
   }
 
-  local line_offset = 0
+  local entries = {}
+  for idx, hunk in pairs(hunks) do
+    table.insert(entries, { idx = idx, hunk = hunk })
+  end
+  table.sort(entries, function(a, b)
+    return a.idx < b.idx
+  end)
 
-  for i, hunk in ipairs(hunks) do
+  for _, entry in ipairs(entries) do
+    local i, hunk = entry.idx, entry.hunk
     local display_hunk = {
       removed_lines = hunk.removed_lines,
       added_lines = hunk.added_lines,
@@ -109,12 +151,19 @@ function M.show(bufnr, hunks, original_lines, proposed_lines)
     end
 
     state.hunks[i] = display_hunk
-    line_offset = line_offset + #hunk.added_lines
   end
 
   buf_states[bufnr] = state
 end
 
+function M.restore(bufnr, snapshot)
+  M.clear(bufnr)
+  if not snapshot then
+    return
+  end
+  M.show(bufnr, snapshot.hunks, snapshot.original_lines, snapshot.proposed_lines)
+end
+
 function M.clear(bufnr)
   vim.api.nvim_buf_clear_namespace(bufnr, ns, 0, -1)
   buf_states[bufnr] = nil