Merge pull request #64 from cposture/fix-visual-selection

fix selection include extra content when the user is in visual mode

4 files changed, 86 insertions, 21 deletions

diff --git a/autoload/vim_ai.vim b/autoload/vim_ai.vim
index 74d6afd..ff447ac 100644
--- a/autoload/vim_ai.vim
+++ b/autoload/vim_ai.vim
@@ -45,27 +45,26 @@ function! vim_ai#MakeScratchWindow()
   endif
 endfunction
 
-function! s:MakeSelectionPrompt(is_selection, lines, instruction, config)
+function! s:MakeSelectionPrompt(selection, instruction, config)
   let l:selection = ""
   if a:instruction == ""
-    let l:selection = a:lines
-  elseif a:is_selection
+    let l:selection = a:selection
+  elseif !empty(a:selection)
     let l:boundary = a:config['options']['selection_boundary']
-    if l:boundary != "" && match(a:lines, l:boundary) == -1
+    if l:boundary != "" && match(a:selection, l:boundary) == -1
       " NOTE: surround selection with boundary (e.g. #####) in order to eliminate empty responses
-      let l:selection = l:boundary . "\n" . a:lines . "\n" . l:boundary
+      let l:selection = l:boundary . "\n" . a:selection . "\n" . l:boundary
     else
-      let l:selection = a:lines
+      let l:selection = a:selection
     endif
   endif
   return l:selection
 endfunction
 
-function! s:MakePrompt(is_selection, lines, instruction, config)
-  let l:lines = trim(join(a:lines, "\n"))
+function! s:MakePrompt(selection, instruction, config)
   let l:instruction = trim(a:instruction)
-  let l:delimiter = l:instruction != "" && a:is_selection ? ":\n" : ""
-  let l:selection = s:MakeSelectionPrompt(a:is_selection, l:lines, l:instruction, a:config)
+  let l:delimiter = l:instruction != "" && a:selection != "" ? ":\n" : ""
+  let l:selection = s:MakeSelectionPrompt(a:selection, l:instruction, a:config)
   return join([l:instruction, l:delimiter, l:selection], "")
 endfunction
 
@@ -88,6 +87,35 @@ function! s:set_nopaste(config)
   endif
 endfunction
 
+function! s:GetCurrentLineOrSelection(is_selection)
+  if a:is_selection
+    return s:GetVisualSelection()
+  else
+    return trim(join(getline(a:firstline, a:lastline), "\n"))
+  endif
+endfunction
+
+function! s:SelectCurrentLineOrSelection(is_selection)
+  if a:is_selection
+    execute "normal! gv"
+  else
+    execute 'normal! V' 
+  endif
+endfunction
+
+function! s:GetVisualSelection()
+  let [line_start, column_start] = getpos("'<")[1:2]
+  let [line_end, column_end] = getpos("'>")[1:2]
+  let lines = getline(line_start, line_end)
+  if len(lines) == 0
+    return ''
+  endif
+  " The exclusive mode means that the last character of the selection area is not included in the operation scope.
+  let lines[-1] = lines[-1][: column_end - (&selection == 'inclusive' ? 1 : 2)]
+  let lines[0] = lines[0][column_start - 1:]
+  return join(lines, "\n")
+endfunction
+
 " Complete prompt
 " - is_selection - <range> parameter
 " - config       - function scoped vim_ai_complete config
@@ -96,15 +124,17 @@ function! vim_ai#AIRun(is_selection, config, ...) range
   let l:config = vim_ai_config#ExtendDeep(g:vim_ai_complete, a:config)
 
   let l:instruction = a:0 ? a:1 : ""
-  let l:lines = getline(a:firstline, a:lastline)
-  let l:prompt = s:MakePrompt(a:is_selection, l:lines, l:instruction, l:config)
+  let l:selection = s:GetCurrentLineOrSelection(a:is_selection)
+  let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
+  " used for getting in Python script
+  let l:is_selection = a:is_selection
 
   let s:last_command = "complete"
   let s:last_config = a:config
   let s:last_instruction = l:instruction
   let s:last_is_selection = a:is_selection
 
-  let l:cursor_on_empty_line = trim(join(l:lines, "\n")) == ""
+  let l:cursor_on_empty_line = empty(getline('.'))
   call s:set_paste(l:config)
   if l:cursor_on_empty_line
     execute "normal! " . a:lastline . "GA"
@@ -124,7 +154,10 @@ function! vim_ai#AIEditRun(is_selection, config, ...) range
   let l:config = vim_ai_config#ExtendDeep(g:vim_ai_edit, a:config)
 
   let l:instruction = a:0 ? a:1 : ""
-  let l:prompt = s:MakePrompt(a:is_selection, getline(a:firstline, a:lastline), l:instruction, l:config)
+  let l:selection = s:GetCurrentLineOrSelection(a:is_selection)
+  " used for getting in Python script
+  let l:is_selection = a:is_selection
+  let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
 
   let s:last_command = "edit"
   let s:last_config = a:config
@@ -132,7 +165,8 @@ function! vim_ai#AIEditRun(is_selection, config, ...) range
   let s:last_is_selection = a:is_selection
 
   call s:set_paste(l:config)
-  execute "normal! " . a:firstline . "GV" . a:lastline . "Gc"
+  call s:SelectCurrentLineOrSelection(a:is_selection)
+  execute "normal! c"
   execute "py3file " . s:complete_py
   call s:set_nopaste(l:config)
 endfunction
@@ -145,7 +179,9 @@ function! vim_ai#AIChatRun(is_selection, config, ...) range
   let l:config = vim_ai_config#ExtendDeep(g:vim_ai_chat, a:config)
 
   let l:instruction = ""
-  let l:lines = getline(a:firstline, a:lastline)
+  let l:selection = s:GetVisualSelection()
+  " used for getting in Python script
+  let l:is_selection = a:is_selection
   call s:set_paste(l:config)
   if &filetype != 'aichat'
     let l:chat_win_id = bufwinid(s:scratch_buffer_name)
@@ -163,7 +199,7 @@ function! vim_ai#AIChatRun(is_selection, config, ...) range
   let l:prompt = ""
   if a:0 || a:is_selection
     let l:instruction = a:0 ? a:1 : ""
-    let l:prompt = s:MakePrompt(a:is_selection, l:lines, l:instruction, l:config)
+    let l:prompt = s:MakePrompt(l:selection, l:instruction, l:config)
   endif
 
   let s:last_command = "chat"
diff --git a/py/chat.py b/py/chat.py
index 6d88015..907e71a 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -52,6 +52,7 @@ initial_messages = parse_chat_messages(initial_prompt)
 
 chat_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
 chat_messages = parse_chat_messages(chat_content)
+is_selection = vim.eval("l:is_selection")
 
 messages = initial_messages + chat_messages
 
@@ -75,7 +76,7 @@ try:
             printDebug("[chat] response: {}", resp)
             return resp['choices'][0]['delta'].get('content', '')
         text_chunks = map(map_chunk, response)
-        render_text_chunks(text_chunks)
+        render_text_chunks(text_chunks, is_selection)
 
         vim.command("normal! a\n\n>>> user\n\n")
         vim.command("redraw")
diff --git a/py/complete.py b/py/complete.py
index 8386c09..ede0aea 100644
--- a/py/complete.py
+++ b/py/complete.py
@@ -9,6 +9,7 @@ openai_options = make_openai_options(config_options)
 http_options = make_http_options(config_options)
 
 prompt = vim.eval("l:prompt").strip()
+is_selection = vim.eval("l:is_selection")
 
 def complete_engine(prompt):
     request = {
@@ -51,7 +52,7 @@ try:
         print('Completing...')
         vim.command("redraw")
         text_chunks = engines[engine](prompt)
-        render_text_chunks(text_chunks)
+        render_text_chunks(text_chunks, is_selection)
         clear_echo_message()
 except BaseException as error:
     handle_completion_error(error)
diff --git a/py/utils.py b/py/utils.py
index 4e638e2..b5ede7f 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -62,14 +62,41 @@ def make_http_options(options):
         'enable_auth': bool(int(options['enable_auth'])),
     }
 
-def render_text_chunks(chunks):
+# During text manipulation in Vim's visual mode, we utilize "normal! c" command. This command deletes the highlighted text,
+# immediately followed by entering insert mode where it generates desirable text.
+
+# Normally, Vim contemplates the position of the first character in selection to decide whether to place the entered text
+# before or after the cursor. For instance, if the given line is "abcd", and "abc" is selected for deletion and "1234" is
+# written in its place, the result is as expected "1234d" rather than "d1234". However, if "bc" is chosen for deletion, the
+# achieved output is "a1234d", whereas "1234ad" is not.
+
+# Despite this, post Vim script's execution of "normal! c", it takes an exit immediately returning to the normal mode. This
+# might trigger a potential misalignment issue especially when the most extreme left character is the line’s second character.
+
+# To avoid such pitfalls, the method "need_insert_before_cursor" checks not only the selection status, but also the character
+# at the first position of the highlighting. If the selection is off or the first position is not the second character in the line,
+# it determines no need for prefixing the cursor.
+def need_insert_before_cursor(is_selection):
+    if is_selection == False:
+        return False
+    pos = vim.eval("getpos(\"'<\")[1:2]")
+    if not isinstance(pos, list) or len(pos) != 2:
+        raise ValueError("Unexpected getpos value, it should be a list with two elements")
+    return pos[1] == "1"
+
+def render_text_chunks(chunks, is_selection):
     generating_text = False
     full_text = ''
+    insert_before_cursor = need_insert_before_cursor(is_selection)
     for text in chunks:
         if not text.strip() and not generating_text:
             continue # trim newlines from the beginning
         generating_text = True
-        vim.command("normal! a" + text)
+        if insert_before_cursor:
+            vim.command("normal! i" + text)
+            insert_before_cursor = False
+        else:
+            vim.command("normal! a" + text)
         vim.command("undojoin")
         vim.command("redraw")
         full_text += text