fix: restore hunk review state on undo and redo

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] = {}

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)

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