feat: Improve diff display and optimize existing prompts for diffs (#134)

deathbeam · web-flow · commit fcace1863096 · 2024-03-08T05:02:58.000+04:00
- Instead of using tooltip just replace the chat window
- Highlight both syntax + diff
- Instead of replacing colors, blend foreground diff colors as background
  colors
- Optimize existing prompts to generate better diffs

Signed-off-by: Tomas Slusny &lt;slusnucky@gmail.com&gt;
diff --git a/README.md b/README.md
@@ -161,6 +161,9 @@ actions.pick(actions.prompt_actions())
 
 - `:CopilotChatExplain` - Explain how it works
 - `:CopilotChatTests` - Briefly explain how selected code works then generate unit tests
+- `:CopilotChatFix` - There is a problem in this code. Rewrite the code to show it with the bug fixed.
+- `:CopilotChatOptimize` - Optimize the selected code to improve performance and readablilty.
+- `:CopilotChatDocs` - 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.
 - `:CopilotChatFixDiagnostic` - Please assist with the following diagnostic issue in file
 - `:CopilotChatCommit` - Write commit message for the change with commitizen convention
 - `:CopilotChatCommitStaged` - Write commit message for the change with commitizen convention
@@ -192,10 +195,19 @@ Also see [here](/lua/CopilotChat/config.lua):
   -- default prompts
   prompts = {
     Explain = {
-      prompt = 'Explain how it works.',
+      prompt = '/COPILOT_EXPLAIN Write a explanation for the code above as paragraphs of text.',
     },
     Tests = {
-      prompt = 'Briefly explain how selected code works then generate unit tests.',
+      prompt = '/COPILOT_TESTS Write a set of detailed unit test functions for the code above.',
+    },
+    Fix = {
+      prompt = '/COPILOT_FIX There is a problem in this code. Rewrite the code to show it with the bug fixed.',
+    },
+    Optimize = {
+      prompt = '/COPILOT_REFACTOR Optimize the selected code to improve performance and readablilty.',
+    },
+    Docs = {
+      prompt = '/COPILOT_REFACTOR 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.',
     },
     FixDiagnostic = {
       prompt = 'Please assist with the following diagnostic issue in file:',
diff --git a/lua/CopilotChat/config.lua b/lua/CopilotChat/config.lua
@@ -79,10 +79,19 @@ return {
   -- default prompts
   prompts = {
     Explain = {
-      prompt = 'Explain how it works.',
+      prompt = '/COPILOT_EXPLAIN Write a explanation for the code above as paragraphs of text.',
     },
     Tests = {
-      prompt = 'Briefly explain how selected code works then generate unit tests.',
+      prompt = '/COPILOT_TESTS Write a set of detailed unit test functions for the code above.',
+    },
+    Fix = {
+      prompt = '/COPILOT_FIX There is a problem in this code. Rewrite the code to show it with the bug fixed.',
+    },
+    Optimize = {
+      prompt = '/COPILOT_REFACTOR Optimize the selected code to improve performance and readablilty.',
+    },
+    Docs = {
+      prompt = '/COPILOT_REFACTOR 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.',
     },
     FixDiagnostic = {
       prompt = 'Please assist with the following diagnostic issue in file:',
diff --git a/lua/CopilotChat/diff.lua b/lua/CopilotChat/diff.lua
@@ -0,0 +1,108 @@
+---@class CopilotChat.Diff
+---@field bufnr number
+---@field current string?
+---@field valid fun(self: CopilotChat.Diff)
+---@field validate fun(self: CopilotChat.Diff)
+---@field show fun(self: CopilotChat.Diff, a: string, b: string, filetype: string, winnr: number)
+---@field restore fun(self: CopilotChat.Diff, winnr: number, bufnr: number)
+
+local utils = require('CopilotChat.utils')
+local is_stable = utils.is_stable
+local class = utils.class
+
+local function blend_color_with_neovim_bg(color_name, blend)
+  local color_int = vim.api.nvim_get_hl(0, { name = color_name }).fg
+  local bg_int = vim.api.nvim_get_hl(0, { name = 'Normal' }).bg
+
+  if not color_int or not bg_int then
+    return
+  end
+
+  local color = { (color_int / 65536) % 256, (color_int / 256) % 256, color_int % 256 }
+  local bg = { (bg_int / 65536) % 256, (bg_int / 256) % 256, bg_int % 256 }
+  local r = math.floor((color[1] * blend + bg[1] * (100 - blend)) / 100)
+  local g = math.floor((color[2] * blend + bg[2] * (100 - blend)) / 100)
+  local b = math.floor((color[3] * blend + bg[3] * (100 - blend)) / 100)
+  return string.format('#%02x%02x%02x', r, g, b)
+end
+
+local function create_buf()
+  local bufnr = vim.api.nvim_create_buf(false, true)
+  vim.api.nvim_buf_set_name(bufnr, 'copilot-diff')
+  vim.bo[bufnr].filetype = 'diff'
+  vim.treesitter.start(bufnr, 'diff')
+  return bufnr
+end
+
+local Diff = class(function(self, on_buf_create, help)
+  self.help = help
+  self.on_buf_create = on_buf_create
+  self.current = nil
+  self.ns = vim.api.nvim_create_namespace('copilot-diff')
+  self.mark_ns = vim.api.nvim_create_namespace('copilot-diff-mark')
+  self.bufnr = create_buf()
+  self.on_buf_create(self.bufnr)
+
+  vim.api.nvim_set_hl(self.ns, '@diff.plus', { bg = blend_color_with_neovim_bg('DiffAdd', 20) })
+  vim.api.nvim_set_hl(self.ns, '@diff.minus', { bg = blend_color_with_neovim_bg('DiffDelete', 20) })
+  vim.api.nvim_set_hl(self.ns, '@diff.delta', { bg = blend_color_with_neovim_bg('DiffChange', 20) })
+end)
+
+function Diff:valid()
+  return self.bufnr and vim.api.nvim_buf_is_valid(self.bufnr)
+end
+
+function Diff:validate()
+  if self:valid() then
+    return
+  end
+  self.bufnr = create_buf()
+  self.on_buf_create(self.bufnr)
+end
+
+function Diff:show(a, b, filetype, winnr)
+  self:validate()
+  self.current = b
+
+  local diff = tostring(vim.diff(a, b, {
+    result_type = 'unified',
+    ignore_blank_lines = true,
+    ignore_whitespace = true,
+    ignore_whitespace_change = true,
+    ignore_whitespace_change_at_eol = true,
+    ignore_cr_at_eol = true,
+    algorithm = 'myers',
+    ctxlen = #a,
+  }))
+  diff = '\n' .. diff
+
+  vim.api.nvim_win_set_buf(winnr, self.bufnr)
+  vim.api.nvim_win_set_hl_ns(winnr, self.ns)
+  vim.bo[self.bufnr].modifiable = true
+  vim.api.nvim_buf_set_lines(self.bufnr, 0, -1, false, vim.split(diff, '\n'))
+  vim.bo[self.bufnr].modifiable = false
+
+  local opts = {
+    id = self.mark_ns,
+    hl_mode = 'combine',
+    priority = 100,
+    virt_text = { { self.help, 'CursorColumn' } },
+  }
+
+  -- stable do not supports virt_text_pos
+  if not is_stable() then
+    opts.virt_text_pos = 'inline'
+  end
+
+  vim.api.nvim_buf_set_extmark(self.bufnr, self.mark_ns, 0, 0, opts)
+  vim.treesitter.start(self.bufnr, 'diff')
+  vim.bo[self.bufnr].syntax = filetype
+end
+
+function Diff:restore(winnr, bufnr)
+  self.current = nil
+  vim.api.nvim_win_set_buf(winnr, bufnr)
+  vim.api.nvim_win_set_hl_ns(winnr, 0)
+end
+
+return Diff
diff --git a/lua/CopilotChat/init.lua b/lua/CopilotChat/init.lua
@@ -2,6 +2,7 @@ local default_config = require('CopilotChat.config')
 local log = require('plenary.log')
 local Copilot = require('CopilotChat.copilot')
 local Chat = require('CopilotChat.chat')
+local Diff = require('CopilotChat.diff')
 local context = require('CopilotChat.context')
 local prompts = require('CopilotChat.prompts')
 local debuginfo = require('CopilotChat.debuginfo')
@@ -13,11 +14,13 @@ local plugin_name = 'CopilotChat.nvim'
 --- @class CopilotChat.state
 --- @field copilot CopilotChat.Copilot?
 --- @field chat CopilotChat.Chat?
+--- @field diff CopilotChat.Diff?
 --- @field source CopilotChat.config.source?
 --- @field config CopilotChat.config?
 local state = {
   copilot = nil,
   chat = nil,
+  diff = nil,
   source = nil,
   config = nil,
 }
@@ -64,21 +67,8 @@ local function show_diff_between_selection_and_copilot(selection)
   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,
-        zindex = M.config.window.zindex + 1,
-      })
-    end
+    local filetype = selection.filetype or vim.bo[state.source.bufnr].filetype
+    state.diff:show(selection.lines, table.concat(lines, '\n'), filetype, state.chat.winnr)
   end
 end
 
@@ -362,7 +352,7 @@ function M.ask(prompt, config, source)
         system_prompt = system_prompt,
         model = config.model,
         temperature = config.temperature,
-        on_done = function(response, token_count)
+        on_done = function(_, token_count)
           if tiktoken.available() and token_count and token_count > 0 then
             append('\n\n' .. token_count .. ' tokens used')
           end
@@ -403,6 +393,47 @@ end
 function M.setup(config)
   M.config = vim.tbl_deep_extend('force', default_config, config or {})
   state.copilot = Copilot(M.config.proxy, M.config.allow_insecure)
+
+  state.diff = Diff(
+    function(bufnr)
+      if M.config.mappings.close then
+        vim.keymap.set('n', M.config.mappings.close, function()
+          state.diff:restore(state.chat.winnr, state.chat.bufnr)
+        end, { buffer = bufnr })
+      end
+      if M.config.mappings.accept_diff then
+        vim.keymap.set('n', M.config.mappings.accept_diff, function()
+          local selection = get_selection()
+          if not selection.start_row or not selection.end_row then
+            return
+          end
+
+          local current = state.diff.current
+          if not current then
+            return
+          end
+
+          local lines = vim.split(current, '\n')
+          if #lines > 0 then
+            vim.api.nvim_buf_set_text(
+              state.source.bufnr,
+              selection.start_row - 1,
+              selection.start_col - 1,
+              selection.end_row - 1,
+              selection.end_col,
+              lines
+            )
+          end
+        end, { buffer = bufnr })
+      end
+    end,
+    "Press '"
+      .. M.config.mappings.close
+      .. "' to close diff, '"
+      .. M.config.mappings.accept_diff
+      .. "' to accept diff."
+  )
+
   state.chat = Chat(plugin_name, function(bufnr)
     if M.config.mappings.complete then
       vim.keymap.set('i', M.config.mappings.complete, complete, { buffer = bufnr })
@@ -413,7 +444,7 @@ function M.setup(config)
     end
 
     if M.config.mappings.close then
-      vim.keymap.set('n', 'q', M.close, { buffer = bufnr })
+      vim.keymap.set('n', M.config.mappings.close, M.close, { buffer = bufnr })
     end
 
     if M.config.mappings.submit_prompt then
diff --git a/lua/CopilotChat/prompts.lua b/lua/CopilotChat/prompts.lua
@@ -33,12 +33,11 @@ The user works in an IDE called Neovim which has a concept for editors with open
 The active document is the source code the user is looking at right now.
 You can only give one reply for each conversation turn.
 You should always generate short suggestions for the next user turns that are relevant to the conversation and not offensive.
-
 ]]
 
 M.COPILOT_EXPLAIN = M.COPILOT_INSTRUCTIONS
   .. [[
-You are an professor of computer science. You are an expert at explaining code to anyone. Your task is to help the Developer understand the code. Pay especially close attention to the selection context.
+You are also an professor of computer science. You are an expert at explaining code to anyone. Your task is to help the Developer understand the code. Pay especially close attention to the selection context.
 
 Additional Rules:
 Provide well thought out examples
@@ -47,53 +46,37 @@ Match the style of provided context when using examples
 Say "I'm not quite sure how to explain that." when you aren't confident in your explanation
 When generating code ensure it's readable and indented properly
 When explaining code, add a final paragraph describing possible ways to improve the code with respect to readability and performance
-
 ]]
 
-M.COPILOT_TESTS = M.COPILOT_INSTRUCTIONS
-  .. [[
-You also specialize in being a highly skilled test generator. Given a description of which test case should be generated, you can generate new test cases. Your task is to help the Developer generate tests. Pay especially close attention to the selection context.
+local preserve_style_rules = [[
 
 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
-don't use private properties or methods from other classes
-Generate the full test file
-Markdown code blocks are used to denote code
+Markdown code blocks are used to denote code.
+If context is provided, try to match the style of the provided code as best as possible. This includes whitespace around the code, at beginning of lines, indentation, and comments.
+Preserve user's code comment blocks, do not exclude them when refactoring code.
+Your code output should keep the same whitespace around the code as the user's 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 in code output as needed to match the user's code.
+Your code output is used for replacing user's code with it so following the above rules is absolutely necessary.
+]]
 
+M.COPILOT_TESTS = M.COPILOT_INSTRUCTIONS
+  .. [[
+You also specialize in being a highly skilled test generator. Given a description of which test case should be generated, you can generate new test cases. Your task is to help the Developer generate tests. Pay especially close attention to the selection context. Do not use private properties or methods from other classes. Generate full test files.
 ]]
+  .. preserve_style_rules
 
 M.COPILOT_FIX = M.COPILOT_INSTRUCTIONS
   .. [[
 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 fix an issue. Pay especially close attention to the selection or exception 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.
-
 ]]
+  .. preserve_style_rules
 
-M.COPILOT_DEVELOPER = M.COPILOT_INSTRUCTIONS
+M.COPILOT_REFACTOR = M.COPILOT_INSTRUCTIONS
   .. [[
 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.)
 ]]
+  .. preserve_style_rules
 
 M.COPILOT_WORKSPACE =
   [[You are a software engineer with expert knowledge of the codebase the user has open in their workspace.
diff --git a/lua/CopilotChat/spinner.lua b/lua/CopilotChat/spinner.lua
@@ -46,7 +46,7 @@ function Spinner:set(text, offset)
       hl_mode = 'combine',
       priority = 100,
       virt_text = vim.tbl_map(function(t)
-        return { t, 'Comment' }
+        return { t, 'CursorColumn' }
       end, vim.split(text, '\n')),
     }
 

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,7 @@ function Spinner:set(text, offset)`
`46`	`46`	`hl_mode = 'combine',`
`47`	`47`	`priority = 100,`
`48`	`48`	`virt_text = vim.tbl_map(function(t)`
`49`		`- return { t, 'Comment' }`
	`49`	`+ return { t, 'CursorColumn' }`
`50`	`50`	`end, vim.split(text, '\n')),`
`51`	`51`	`}`
`52`	`52`