feat(CopilotChat): Improve line separation and prompt update logic, add diff view (#99)

* feat(CopilotChat): Improve line separation and prompt update logic

Now submit_code and submit_prompt properly try to find first match from the end
so you dont need to cursor over specific blocks. This is also prerequisite for
showing current diff as float or in corner somewhere.

BREAKING CHANGE: The key mapping for submitting code is now <C-d> instead
of <C-y>.

Signed-off-by: Tomas Slusny <ts6234@att.com>

* Improve the default prompts slightly and fix issue with active selection

Signed-off-by: Tomas Slusny <ts6234@att.com>

* feat: Add diff display

For now with keybinding (default K). Maybe find a way to display it in better place.

Signed-off-by: Tomas Slusny <ts6234@att.com>

* Adjust keybindings for accept_diff and show_diff

---------

Signed-off-by: Tomas Slusny <ts6234@att.com>
Co-authored-by: Tomas Slusny <ts6234@att.com>

Author: deathbeam

lua/CopilotChat/copilot.lua

@@ -70,7 +70,7 @@ local function generate_request(history, selection, filetype, system_prompt, mod
   if selection ~= '' then
     -- Insert the active selection before last prompt
     table.insert(messages, #messages, {
-      content = '\nActive selection:\n```' .. selection .. '\n' .. filetype .. '\n```',
+      content = '\nActive selection:\n```' .. filetype .. '\n' .. selection .. '\n```',
       role = 'system',
     })
   end

lua/CopilotChat/init.lua

@@ -23,53 +23,88 @@ function CopilotChatFoldExpr(lnum, separator)
   return '='
 end
 
-local function find_lines_between_separator_at_cursor(bufnr, separator)
-  local lines = vim.api.nvim_buf_get_lines(bufnr, 0, -1, false)
-  local cursor = vim.api.nvim_win_get_cursor(0)
-  local cursor_line = cursor[1]
+local function find_lines_between_separator(lines, pattern, at_least_one)
   local line_count = #lines
-  local last_separator_line = 1
-  local next_separator_line = line_count
-  local pattern = '^' .. separator .. '%w*$'
+  local separator_line_start = 1
+  local separator_line_finish = line_count
+  local found_one = false
 
   -- Find the last occurrence of the separator
-  for i, line in ipairs(lines) do
-    if i > cursor_line and string.find(line, pattern) then
-      next_separator_line = i - 1
-      break
-    end
+  for i = line_count, 1, -1 do -- Reverse the loop to start from the end
+    local line = lines[i]
     if string.find(line, pattern) then
-      last_separator_line = i + 1
+      if i < (separator_line_finish + 1) and (not at_least_one or found_one) then
+        separator_line_start = i + 1
+        break -- Exit the loop as soon as the condition is met
+      end
+
+      found_one = true
+      separator_line_finish = i - 1
     end
   end
 
+  if at_least_one and not found_one then
+    return {}, 1, 1, 0
+  end
+
   -- Extract everything between the last and next separator
   local result = {}
-  for i = last_separator_line, next_separator_line do
+  for i = separator_line_start, separator_line_finish do
     table.insert(result, lines[i])
   end
 
-  return table.concat(result, '\n'), last_separator_line, next_separator_line, line_count
+  return result, separator_line_start, separator_line_finish, line_count
+end
+
+local function show_diff_between_selection_and_copilot()
+  local selection = state.selection
+  if not selection or not selection.buffer or not selection.start_row or not selection.end_row then
+    return
+  end
+
+  local chat_lines = vim.api.nvim_buf_get_lines(state.chat.bufnr, 0, -1, false)
+  local section_lines = find_lines_between_separator(chat_lines, M.config.separator .. '$', true)
+  local lines = find_lines_between_separator(section_lines, '^```%w*$', true)
+  if #lines > 0 then
+    local diff = tostring(vim.diff(selection.lines, table.concat(lines, '\n'), {}))
+    if diff and diff ~= '' then
+      vim.lsp.util.open_floating_preview(vim.split(diff, '\n'), 'diff', {
+        border = 'single',
+        title = M.config.name .. ' Diff',
+        title_pos = 'left',
+        focusable = false,
+        focus = false,
+        relative = 'editor',
+        row = 0,
+        col = 0,
+        width = vim.api.nvim_win_get_width(0) - 3,
+      })
+    end
+  end
 end
 
 local function update_prompts(prompt, system_prompt)
   local prompts_to_use = M.get_prompts()
+  local try_again = false
   local result = string.gsub(prompt, [[/[%w_]+]], function(match)
-    match = string.sub(match, 2)
-    local found = prompts_to_use[match]
-
+    local found = prompts_to_use[string.sub(match, 2)]
     if found then
       if found.kind == 'user' then
-        return found.prompt
+        local out = found.prompt
+        if string.match(out, [[/[%w_]+]]) then
+          try_again = true
+        end
+        return out
       elseif found.kind == 'system' then
         system_prompt = found.prompt
+        return ''
       end
     end
 
-    return ''
+    return match
   end)
 
-  if string.match(result, [[/[%w_]+]]) then
+  if try_again then
     return update_prompts(result, system_prompt)
   end
 
@@ -204,9 +239,10 @@ function M.open(config)
 
     if config.mappings.submit_prompt then
       vim.keymap.set('n', config.mappings.submit_prompt, function()
-        local input, start_line, end_line, line_count =
-          find_lines_between_separator_at_cursor(state.chat.bufnr, config.separator)
-        input = vim.trim(input)
+        local chat_lines = vim.api.nvim_buf_get_lines(state.chat.bufnr, 0, -1, false)
+        local lines, start_line, end_line, line_count =
+          find_lines_between_separator(chat_lines, config.separator .. '$')
+        local input = vim.trim(table.concat(lines, '\n'))
         if input ~= '' then
           -- If we are entering the input at the end, replace it
           if line_count == end_line then
@@ -217,8 +253,14 @@ function M.open(config)
       end, { buffer = state.chat.bufnr })
     end
 
-    if config.mappings.submit_code then
-      vim.keymap.set('n', config.mappings.submit_code, function()
+    if config.mappings.show_diff then
+      vim.keymap.set('n', config.mappings.show_diff, show_diff_between_selection_and_copilot, {
+        buffer = state.chat.bufnr,
+      })
+    end
+
+    if config.mappings.accept_diff then
+      vim.keymap.set('n', config.mappings.accept_diff, function()
         if
           not state.selection
           or not state.selection.buffer
@@ -229,15 +271,18 @@ function M.open(config)
           return
         end
 
-        local input = find_lines_between_separator_at_cursor(state.chat.bufnr, '```')
-        if input ~= '' then
+        local chat_lines = vim.api.nvim_buf_get_lines(state.chat.bufnr, 0, -1, false)
+        local section_lines =
+          find_lines_between_separator(chat_lines, config.separator .. '$', true)
+        local lines = find_lines_between_separator(section_lines, '^```%w*$', true)
+        if #lines > 0 then
           vim.api.nvim_buf_set_text(
             state.selection.buffer,
             state.selection.start_row - 1,
             state.selection.start_col - 1,
             state.selection.end_row - 1,
             state.selection.end_col,
-            vim.split(input, '\n')
+            lines
           )
         end
       end, { buffer = state.chat.bufnr })
@@ -261,6 +306,7 @@ function M.open(config)
     elseif layout == 'horizontal' then
       win_opts.vertical = false
     elseif layout == 'float' then
+      win_opts.zindex = 1
       win_opts.relative = config.window.relative
       win_opts.border = config.window.border
       win_opts.title = config.window.title
@@ -466,7 +512,8 @@ M.config = {
     reset = '<C-l>',
     complete = '<Tab>',
     submit_prompt = '<CR>',
-    submit_code = '<C-y>',
+    accept_diff = '<C-y>',
+    show_diff = '<C-d>',
   },
 }
 --- Set up the plugin

lua/CopilotChat/prompts.lua

@@ -75,15 +75,25 @@ Preserve user's code comment blocks, do not exclude them when refactoring code.
 
 M.COPILOT_DEVELOPER = M.COPILOT_INSTRUCTIONS
   .. [[
-You're a 10x senior developer that is an expert in programming.
-Your job is to change the user's code according to their needs.
-Your job is only to change / edit the code.
-Your code output should keep the same level of indentation as the user's code.
-You MUST add whitespace in the beginning of each line as needed to match the user's code.
+You also specialize in being a highly skilled code generator. Given a description of what to do you can refactor, modify or enhance existing code. Your task is help the Developer change their code according to their needs. Pay especially close attention to the selection context.
 
+Additional Rules:
+If context is provided, try to match the style of the provided code as best as possible
+Generated code is readable and properly indented
+Markdown blocks are used to denote code
+Preserve user's code comment blocks, do not exclude them when refactoring code.
+
+]]
+
+M.USER_EXPLAIN = 'Write a explanation for the code above as paragraphs of text.'
+M.USER_TESTS = 'Write a set of detailed unit test functions for the code above.'
+M.USER_FIX = 'There is a problem in this code. Rewrite the code to show it with the bug fixed.'
+M.USER_DOCS = [[Write documentation for the selected code.
+The reply should be a codeblock containing the original code with the documentation added as comments.
+Use the most appropriate documentation style for the programming language used (e.g. JSDoc for JavaScript, docstrings for Python etc.)
 ]]
 
-M.COPILOT_WORKSPACE =
+COPILOT_WORKSPACE =
   [[You are a software engineer with expert knowledge of the codebase the user has open in their workspace.
 When asked for your name, you must respond with "GitHub Copilot".
 Follow the user's requirements carefully & to the letter.
@@ -144,14 +154,6 @@ Response:
 To read a file, you can use a [`FileReader`](src/fs/fileReader.ts) class from [src/fs/fileReader.ts](src/fs/fileReader.ts).
 ]]
 
-M.USER_EXPLAIN = 'Write a explanation for the code above as paragraphs of text.'
-M.USER_TESTS = 'Write a set of detailed unit test functions for the code above.'
-M.USER_FIX = 'There is a problem in this code. Rewrite the code to show it with the bug fixed.'
-M.USER_DOCS = [[Write documentation for the selected code.
-The reply should be a codeblock containing the original code with the documentation added as comments.
-Use the most appropriate documentation style for the programming language used (e.g. JSDoc for JavaScript, docstrings for Python etc.)
-]]
-
 EMBEDDING_KEYWORDS =
   [[You are a coding assistant who help the user answer questions about code in their workspace by providing a list of relevant keywords they can search for to answer the question.
 The user will provide you with potentially relevant information from the workspace. This information may be incomplete.