From 40b53dea83fae73e90236a5f82cdd3c2670834ff Mon Sep 17 00:00:00 2001
From: danymat <d.danymat@gmail.com>
Date: Mon, 10 Jan 2022 17:23:53 +0100
Subject: [PATCH] feat: Return to previous position after las cycling (closes
 #35)

---
 lua/neogen.lua                  |  8 ++++++--
 lua/neogen/utilities/cursor.lua | 23 +++++++++--------------
 2 files changed, 15 insertions(+), 16 deletions(-)

diff --git a/lua/neogen.lua b/lua/neogen.lua
index 896b9d8..7ed0b5f 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -91,7 +91,7 @@ neogen.generate = function(opts)
                 -- First and last extmarks are needed to know the range of inserted content
                 if neogen.configuration.input_after_comment == true then
                     -- Creates extmark for the beggining of the content
-                    neogen.utilities.cursor.create(to_place + 1, start_column, true)
+                    neogen.utilities.cursor.create(to_place + 1, start_column)
                     -- Creates extmarks for the content
                     for i, value in pairs(content_with_marks) do
                         local start = 0
@@ -106,8 +106,12 @@ neogen.generate = function(opts)
                         end
                     end
 
+                    -- Create extmark to jump back to current location
+                    local pos = vim.api.nvim_win_get_cursor(0)
+                    neogen.utilities.cursor.create(pos[1], pos[2])
+
                     -- Creates extmark for the end of the content
-                    neogen.utilities.cursor.create(to_place + #content + 1, 0, true)
+                    neogen.utilities.cursor.create(to_place + #content + 1, 0)
 
                     neogen.utilities.cursor.jump({ first_time = true })
                 end
diff --git a/lua/neogen/utilities/cursor.lua b/lua/neogen/utilities/cursor.lua
index 81f54a6..1dd6257 100644
--- a/lua/neogen/utilities/cursor.lua
+++ b/lua/neogen/utilities/cursor.lua
@@ -5,15 +5,12 @@ local neogen_virt_text_ns = vim.api.nvim_create_namespace("neogen_virt_text")
 local current_position = 1
 
 --- Wrapper around set_extmark with 1-based numbering for `line` and `col`, and returns the id of the created extmark
---- @param line string
---- @param col string
+--- @param line number
+--- @param col number
 --- @return number
-neogen.utilities.cursor.create = function(line, col, delimiters)
-    local opts = delimiters == true and {}
-        or { virt_text = { { "TODO", "TSTitle" } }, virt_text_pos = "overlay", virt_text_hide = true }
+neogen.utilities.cursor.create = function(line, col)
     current_position = 1
     local new_col = col == 0 and 0 or col - 1
-    -- vim.api.nvim_buf_set_extmark(0, neogen_virt_text_ns, line - 1, new_col, opts)
     return vim.api.nvim_buf_set_extmark(0, neogen_ns, line - 1, new_col, {})
 end
 
@@ -22,18 +19,13 @@ end
 --- First jumpable extmark is the one after the extmarks responsible of start/end of annotation
 neogen.utilities.cursor.go_next_extmark = function()
     local extm_list = vim.api.nvim_buf_get_extmarks(0, neogen_ns, 0, -1, {})
-    local virt_text_extm = vim.api.nvim_buf_get_extmarks(0, neogen_virt_text_ns, 0, -1, {})
     local position = current_position + 1
+    table.sort(extm_list, function(a, b)
+        return a[1] < b[1]
+    end)
 
     if #extm_list ~= 2 then
         local pos = { extm_list[position][2] + 1, extm_list[position][3] }
-        local virt_text = vim.tbl_filter(function(e)
-            return e[2] == extm_list[position][2] and e[3] == extm_list[position][3]
-        end, virt_text_extm)[1]
-        if virt_text and #virt_text > 0 then
-            vim.api.nvim_buf_del_extmark(0, neogen_virt_text_ns, virt_text[1])
-        end
-
         vim.api.nvim_win_set_cursor(0, pos)
         current_position = current_position + 1
         return true
@@ -61,6 +53,9 @@ end
 
 neogen.utilities.cursor.jump_prev = function()
     local marks = vim.api.nvim_buf_get_extmarks(0, neogen_ns, 0, -1, {})
+    table.sort(marks, function(a, b)
+        return a[1] < b[1]
+    end)
 
     if #marks == 2 then
         return false