From a4ce0f24e33e0c743d6a6e2c001560000bd1fc9e Mon Sep 17 00:00:00 2001
From: Ajay Mamtora <45173937+Ajaymamtora@users.noreply.github.com>
Date: Tue, 13 Aug 2024 21:43:28 +0100
Subject: [PATCH] feat(action): Allow passing a range to filter the undo items
 by ? #57

---
 lua/telescope-undo/init.lua | 43 ++++++++++++++++++++++---------------
 1 file changed, 26 insertions(+), 17 deletions(-)

diff --git a/lua/telescope-undo/init.lua b/lua/telescope-undo/init.lua
index 54f750c..e6659ea 100644
--- a/lua/telescope-undo/init.lua
+++ b/lua/telescope-undo/init.lua
@@ -6,7 +6,7 @@ require("telescope-undo.actions")
 local get_previewer = require("telescope-undo.previewer").get_previewer
 local timeago = require("telescope-undo.lua-timeago").timeago
 
-local function _traverse_undotree(opts, entries, level)
+local function _traverse_undotree(opts, entries, level, line_range)
   local undolist = {}
   -- create diffs for each entry in our undotree
   for i = #entries, 1, -1 do
@@ -15,12 +15,12 @@ local function _traverse_undotree(opts, entries, level)
     end
     -- grab the buffer as it is after this iteration's undo state
     vim.cmd("silent undo " .. entries[i].seq)
-    local buffer_after_lines = vim.api.nvim_buf_get_lines(0, 0, -1, false) or {}
+    local buffer_after_lines = vim.api.nvim_buf_get_lines(0, line_range[1] - 1, line_range[2], false) or {}
     local buffer_after = table.concat(buffer_after_lines, "\n")
 
     -- grab the buffer as it is after this undo state's parent
     vim.cmd("silent undo")
-    local buffer_before_lines = vim.api.nvim_buf_get_lines(0, 0, -1, false) or {}
+    local buffer_before_lines = vim.api.nvim_buf_get_lines(0, line_range[1] - 1, line_range[2], false) or {}
     local buffer_before = table.concat(buffer_before_lines, "\n")
 
     -- build diff header so that delta can go ahead and syntax highlight
@@ -46,22 +46,26 @@ local function _traverse_undotree(opts, entries, level)
       end
     end
 
-    -- use the data we just created to feed into our finder later
-    table.insert(undolist, {
-      seq = entries[i].seq, -- save state number, used in display and to restore
-      alt = level, -- current level, i.e. how deep into alt branches are we, used to graph
-      first = i == #entries, -- whether this is the first node in this branch, used to graph
-      time = entries[i].time, -- save state time, used in display
-      ordinal = ordinal, -- a long string of all additions and deletions, used for search
-      diff = header .. diff, -- the proper diff, used for preview
-      additions = additions, -- all additions, used to yank a result
-      deletions = deletions, -- all deletions, used to yank a result
-      bufnr = vim.api.nvim_get_current_buf(), -- for which buffer this telescope was invoked, used to restore
-    })
+    -- Only add to undolist if there are changes within the specified range
+    if diff ~= "" then
+      -- use the data we just created to feed into our finder later
+      table.insert(undolist, {
+        seq = entries[i].seq, -- save state number, used in display and to restore
+        alt = level, -- current level, i.e. how deep into alt branches are we, used to graph
+        first = i == #entries, -- whether this is the first node in this branch, used to graph
+        time = entries[i].time, -- save state time, used in display
+        ordinal = ordinal, -- a long string of all additions and deletions, used for search
+        diff = header .. diff, -- the proper diff, used for preview
+        additions = additions, -- all additions, used to yank a result
+        deletions = deletions, -- all deletions, used to yank a result
+        bufnr = vim.api.nvim_get_current_buf(), -- for which buffer this telescope was invoked, used to restore
+        line_range = line_range, -- Add line range information
+      })
+    end
 
     -- descend recursively into alternate histories of undo states
     if entries[i].alt ~= nil then
-      local alt_undolist = _traverse_undotree(opts, entries[i].alt, level + 1)
+      local alt_undolist = _traverse_undotree(opts, entries[i].alt, level + 1, line_range)
       -- pretend these results are our results
       for _, elem in pairs(alt_undolist) do
         table.insert(undolist, elem)
@@ -80,7 +84,8 @@ local function build_undolist(opts)
   local ut = vim.fn.undotree()
 
   -- TODO: maybe use this opportunity to limit the number of root nodes we process overall, to ensure good performance
-  local undolist = _traverse_undotree(opts, ut.entries, 0)
+  -- Pass the line range to _traverse_undotree
+  local undolist = _traverse_undotree(opts, ut.entries, 0, opts.line_range)
 
   -- restore everything after all diffs have been created
   -- BUG: `gi` (last insert location) is being killed by our method, we should save that as well
@@ -98,6 +103,10 @@ M.undo = function(opts)
     return
   end
   opts = opts or {}
+
+  -- Add line range parameter
+  opts.line_range = opts.line_range or { 1, -1 } -- Default to entire buffer
+
   pickers
     .new(opts, {
       prompt_title = "Undo History",