fix: patch #2133 and handle more variants of FIM suggestions (#2407)

* fix: patch #2133 and handle more variants of FIM suggestions

* fix: only consider FIM if the first edit line matches trigger position

* fix: test

---------

Co-authored-by: atontb <[email protected]>

Author: Will-ShaoHua

server/aws-lsp-codewhisperer/src/language-server/inline-completion/handler/editCompletionHandler.ts

@@ -400,7 +400,8 @@ export class EditCompletionHandler {
                     const processedSuggestion = processEditSuggestion(
                         suggestion.content,
                         session.startPosition,
-                        session.document
+                        session.document,
+                        session.requestContext.fileContext.rightFileContent
                     )
                     const isInlineEdit = processedSuggestion.type === SuggestionType.EDIT
 

server/aws-lsp-codewhisperer/src/language-server/inline-completion/utils/diffUtils.test.ts

@@ -123,6 +123,27 @@ describe('categorizeUnifieddiffV2v2 should return correct type (addOnly, edit, d
 
     describe('addOnly', function () {
         const addOnlyCases: Case[] = [
+            {
+                udiff: `--- a/src/main/hello/MathUtil.java
++++ b/src/main/hello/MathUtil.java
+@@ -10,6 +10,6 @@ public class MathUtil {
+     }   
+ 
+     public static int multiply(int a, int b) {
+-        return a * b;
++        return a * b * c;
+     }
+ }`,
+            },
+            {
+                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
++++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
+@@ -6,2 +6,3 @@
+-        return a * b;
++        return a * b *
++               c * d;
+ }`,
+            },
             {
                 udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
 +++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
@@ -141,12 +162,10 @@ describe('categorizeUnifieddiffV2v2 should return correct type (addOnly, edit, d
  }`,
             },
             {
-                udiff: `--- file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
-+++ file:///Volumes/workplace/ide/sample_projects/Calculator-2/src/main/hello/MathUtil.java
-@@ -6,7 +6,11 @@
- 
-     // write a function to subtract 2 numbers
-      public static int subtract(int a, int b) {
+                udiff: `--- a/src/main/hello/MathUtil.java
++++ b/src/main/hello/MathUtil.java
+@@ -8,5 +8,10 @@ public class MathUtil {
+     public static int subtract(int a, int b) {
          return a - b;
      }
 -    
@@ -251,20 +270,14 @@ describe('categorizeUnifieddiffV2v2 should return correct type (addOnly, edit, d
             {
                 udiff: `--- a/src/main/hello/MathUtil.java
 +++ b/src/main/hello/MathUtil.java
-@@ -1,11 +1,11 @@
+@@ -1,6 +1,6 @@
  public class MathUtil {
      // write a function to add 2 numbers
 -    public static int add(int a, int b) {
 +    public static double add(double a, double b) {
          return a + b;
      }
- 
-     // write a function to subtract 2 numbers
-     public static int subtract(int a, int b) {
-     public static double subtract(double a, double b) {
-         return a - b;
-     }   
- }`,
+ `,
             },
             {
                 udiff: `--- a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.ts
@@ -300,6 +313,17 @@ describe('categorizeUnifieddiffV2v2 should return correct type (addOnly, edit, d
              const languageId = getLanguageIdFromUri(uri)
              textDocument = TextDocument.create(uri, languageId, 0, content)`,
             },
+            {
+                udiff: `--- a/src/main/hello/MathUtil.java
++++ b/src/main/hello/MathUtil.java
+@@ -12,4 +12,5 @@ public class MathUtil {
+ 
+     // write a function to multiply 2 numbers
+     public static int multiply(int a, int b) {
+-        return a * b;
++        return a * b * c;
++    }`,
+            },
         ]
 
         for (let i = 0; i < cases.length; i++) {

server/aws-lsp-codewhisperer/src/language-server/inline-completion/utils/diffUtils.ts

@@ -3,12 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import * as diff from 'diff'
+import { parsePatch, Hunk, createTwoFilesPatch } from 'diff'
 import { CodeWhispererSupplementalContext, CodeWhispererSupplementalContextItem } from '../../../shared/models/model'
 import { trimSupplementalContexts } from '../../../shared/supplementalContextUtil/supplementalContextUtil'
 import { Position, TextDocument, Range } from '@aws/language-server-runtimes/protocol'
 import { SuggestionType } from '../../../shared/codeWhispererService'
-import { getPrefixSuffixOverlap } from './mergeRightUtils'
+import { getPrefixSuffixOverlap, truncateOverlapWithRightContext } from './mergeRightUtils'
 
 /**
  * Generates a unified diff format between old and new file contents
@@ -31,7 +31,7 @@ export function generateUnifiedDiffWithTimestamps(
     newTimestamp: number,
     contextSize: number = 3
 ): string {
-    const patchResult = diff.createTwoFilesPatch(
+    const patchResult = createTwoFilesPatch(
         oldFilePath,
         newFilePath,
         oldContent,
@@ -204,12 +204,13 @@ export function getCharacterDifferences(addedLines: string[], deletedLines: stri
 export function processEditSuggestion(
     unifiedDiff: string,
     triggerPosition: Position,
-    document: TextDocument
+    document: TextDocument,
+    rightContext: string
 ): { suggestionContent: string; type: SuggestionType } {
     // Assume it's an edit if anything goes wrong, at the very least it will not be rendered incorrectly
     let diffCategory: ReturnType<typeof categorizeUnifieddiff> = 'edit'
     try {
-        diffCategory = categorizeUnifieddiff(unifiedDiff)
+        diffCategory = categorizeUnifieddiff(unifiedDiff, triggerPosition.line)
     } catch (e) {
         // We dont have logger here....
         diffCategory = 'edit'
@@ -228,8 +229,9 @@ export function processEditSuggestion(
          * if LSP returns `g('foo')` instead of `.log()` the suggestion will be discarded because prefix doesnt match
          */
         const processedAdd = removeOverlapCodeFromSuggestion(leftContextAtTriggerLine, preprocessAdd)
+        const mergedWithRightContext = truncateOverlapWithRightContext(rightContext, processedAdd)
         return {
-            suggestionContent: processedAdd,
+            suggestionContent: mergedWithRightContext,
             type: SuggestionType.COMPLETION,
         }
     } else {
@@ -247,10 +249,26 @@ interface UnifiedDiff {
     firstPlusIndex: number
     minusIndexes: number[]
     plusIndexes: number[]
+    hunk: Hunk
 }
 
 // TODO: refine
 export function readUdiff(unifiedDiff: string): UnifiedDiff {
+    let hunk: Hunk | undefined
+    try {
+        const patches = parsePatch(unifiedDiff)
+        if (patches.length !== 1) {
+            throw new Error(`Provided unified diff from has 0 or more than 1 patches`)
+        }
+        hunk = patches[0].hunks[0]
+        if (!hunk) {
+            throw new Error(`Null hunk`)
+        }
+    } catch (e) {
+        throw e
+    }
+
+    // TODO: Should use hunk instead of parsing manually
     const lines = unifiedDiff.split('\n')
     const headerEndIndex = lines.findIndex(l => l.startsWith('@@'))
     if (headerEndIndex === -1) {
@@ -275,30 +293,64 @@ export function readUdiff(unifiedDiff: string): UnifiedDiff {
     const firstMinusIndex = relevantLines.findIndex(s => s.startsWith('-'))
     const firstPlusIndex = relevantLines.findIndex(s => s.startsWith('+'))
 
+    // TODO: Comment these out as they are used for a different version of addonly type determination logic in case the current implementation doesn't work.
+    // Could remove later if we are sure current imple works.
+    /**
+     * Concatenate all contiguous added lines (i.e., unbroken sequence of "+"s).
+     * Exclude all newlines when concatenating, so we get a single line representing the new text
+     */
+    // let singleLine = ''
+    // let prev: number | undefined
+    // for (const idx of plusIndexes) {
+    //     if (!prev || idx === prev + 1) {
+    //         const removedPlus = relevantLines[idx].substring(1)
+    //         const removedStartNewline = trimStartNewline(removedPlus)
+    //         singleLine += removedStartNewline
+    //     } else {
+    //         break
+    //     }
+    // }
+
     return {
         linesWithoutHeaders: relevantLines,
         firstMinusIndex: firstMinusIndex,
         firstPlusIndex: firstPlusIndex,
         minusIndexes: minusIndexes,
         plusIndexes: plusIndexes,
+        hunk: hunk,
     }
 }
 
-export function categorizeUnifieddiff(unifiedDiff: string): 'addOnly' | 'deleteOnly' | 'edit' {
+// Theoretically, we should always pass userTriggerAtLine, keeping it nullable for easier testing for now
+export function categorizeUnifieddiff(
+    unifiedDiff: string,
+    userTriggerAtLine?: number
+): 'addOnly' | 'deleteOnly' | 'edit' {
     try {
         const d = readUdiff(unifiedDiff)
+        const hunk = d.hunk
         const firstMinusIndex = d.firstMinusIndex
         const firstPlusIndex = d.firstPlusIndex
         const diffWithoutHeaders = d.linesWithoutHeaders
 
+        // Shouldn't be the case but if there is no - nor +, assume it's an edit
         if (firstMinusIndex === -1 && firstPlusIndex === -1) {
             return 'edit'
         }
 
+        // If first "EDIT" line is not where users trigger, it must be EDIT
+        // Note hunk.start is 1 based index
+        const firstLineEdited = hunk.oldStart - 1 + Math.min(...d.minusIndexes, ...d.plusIndexes)
+        if (userTriggerAtLine !== undefined && userTriggerAtLine !== firstLineEdited) {
+            return 'edit'
+        }
+
+        // Naive case, only +
         if (firstMinusIndex === -1 && firstPlusIndex !== -1) {
             return 'addOnly'
         }
 
+        // Naive case, only -
         if (firstMinusIndex !== -1 && firstPlusIndex === -1) {
             return 'deleteOnly'
         }
@@ -321,12 +373,47 @@ export function categorizeUnifieddiff(unifiedDiff: string): 'addOnly' | 'deleteO
 
         // If last '-' line is followed by '+' block, it could be addonly
         if (plusIndexes[0] === minusIndexes[minusIndexes.length - 1] + 1) {
+            /**
+            -------------------------------
+            -  return 
+            +  return a - b;
+            -------------------------------
+            commonPrefix = "return "
+            minusLinesDelta = ""
+
+            --------------------------------
+            -\t\t\t
+            +\treturn a - b;
+            --------------------------------
+            commonPrefix = "\t"
+            minusLinesDelta = "\t\t"
+
+             * 
+             * 
+             * 
+             */
             const minusLine = diffWithoutHeaders[minusIndexes[minusIndexes.length - 1]].substring(1)
             const pluscode = extractAdditions(unifiedDiff)
 
             // If minusLine subtract the longest common substring of minusLine and plugcode and it's empty string, it's addonly
             const commonPrefix = longestCommonPrefix(minusLine, pluscode)
-            if (minusLine.substring(commonPrefix.length).trim().length === 0) {
+            const minusLinesDelta = minusLine.substring(commonPrefix.length)
+            if (minusLinesDelta.trim().length === 0) {
+                return 'addOnly'
+            }
+
+            /**
+            -------------------------------
+             -  return a * b;
+             +  return a * b * c;
+            -------------------------------
+            commonPrefix = "return a * b"
+            minusLinesDelta = ";"
+            pluscodeDelta = " * c;"
+             *
+             */
+            const pluscodeDelta = pluscode.substring(commonPrefix.length)
+            if (pluscodeDelta.endsWith(minusLinesDelta)) {
                 return 'addOnly'
             }
         }
@@ -400,3 +487,23 @@ export function longestCommonPrefix(str1: string, str2: string): string {
 
     return prefix
 }
+
+// TODO: They are used for a different version of addonly type determination logic in case the current implementation doesn't work.
+// Could remove later if we are sure current impl works.
+// function trimStartNewline(str: string): string {
+//     return str.replace(/^[\n\r]+/, '')
+// }
+
+// function hasOneContiguousInsert(original: string, changed: string) {
+//     const delta = changed.length - original.length
+//     if (delta <= 0) {
+//         // Changed string must be longer
+//         return false
+//     }
+
+//     let p, s
+//     for (p = 0; original[p] === changed[p] && p < original.length; ++p);
+//     for (s = original.length - 1; original[s] === changed[s + delta] && s >= 0; --s);
+
+//     return p === s + 1
+// }