Fix formatting for single line completions

gradle.properties

@@ -13,7 +13,7 @@ platformType=IC
 platformVersion=2022.1
 # Plugin Dependencies -> https://plugins.jetbrains.com/docs/intellij/plugin-dependencies.html
 # Example: platformPlugins = com.intellij.java, com.jetbrains.php:203.4449.22
-platformPlugins=Git4Idea,PerforceDirectPlugin
+platformPlugins=Git4Idea,PerforceDirectPlugin,java
 # Java language level used to compile sources and to generate the files for - Java 11 is required for 2020.3 <= x < 2022.2
 javaVersion=11
 # Gradle Releases -> https://github.com/gradle/gradle/releases

src/main/kotlin/com/sourcegraph/cody/autocomplete/CodyAutocompleteManager.kt

@@ -329,32 +329,24 @@ class CodyAutocompleteManager {
       triggerKind: InlineCompletionTriggerKind,
   ) {
     val project = editor.project
-    if (project != null && System.getProperty("cody.autocomplete.enableFormatting") != "false") {
-      items.map { item ->
-        if (item.insertText.lines().size > 1) {
-          item.insertText =
-              item.insertText.lines()[0] +
-                  CodyFormatter.formatStringBasedOnDocument(
-                      item.insertText.lines().drop(1).joinToString(separator = "\n"),
-                      project,
-                      editor.document,
-                      offset)
-        }
-      }
-    }
-
     val defaultItem = items.firstOrNull() ?: return
     val range = getTextRange(editor.document, defaultItem.range)
     val originalText = editor.document.getText(range)
-    val lines = defaultItem.insertText.lines()
-    val insertTextFirstLine: String = lines.firstOrNull() ?: ""
-    val multilineInsertText: String = lines.drop(1).joinToString(separator = "\n")
 
-    // Run Myers diff between the existing text in the document and the first line of the
-    // `insertText` that is returned from the agent.
+    val formattedInsertText =
+        if (project == null ||
+            System.getProperty("cody.autocomplete.enableFormatting") == "false") {
+          defaultItem.insertText
+        } else {
+          CodyFormatter.formatStringBasedOnDocument(
+              defaultItem.insertText, project, editor.document, range, offset)
+        }
+
+    // Run Myers diff between the existing text in the document and the `insertText` that is
+    // returned from the agent.
     // The diff algorithm returns a list of "deltas" that give us the minimal number of additions we
     // need to make to the document.
-    val patch = diff(originalText, insertTextFirstLine)
+    val patch = diff(originalText, defaultItem.insertText)
     if (!patch.deltas.all { delta -> delta.type == Delta.TYPE.INSERT }) {
       if (triggerKind == InlineCompletionTriggerKind.INVOKE ||
           UserLevelConfig.isVerboseLoggingEnabled()) {
@@ -371,26 +363,29 @@ class CodyAutocompleteManager {
       }
     }
 
-    // Insert one inlay hint per delta in the first line.
-    for (delta in patch.deltas) {
-      val text = delta.revised.lines.joinToString("")
-      inlayModel.addInlineElement(
-          range.startOffset + delta.original.position,
-          true,
-          CodyAutocompleteSingleLineRenderer(text, items, editor, AutocompleteRendererType.INLINE))
-    }
+    defaultItem.insertText = formattedInsertText
+    val completionText = formattedInsertText.removePrefix(originalText)
 
-    // Insert remaining lines of multiline completions as a single block element under the
-    // (potentially false?) assumption that we don't need to compute diffs for them. My
-    // understanding of multiline completions is that they are only supposed to be triggered in
-    // situations where we insert a large block of code in an empty block.
-    if (multilineInsertText.isNotEmpty()) {
+    val lineBreaks = listOf("\r\n", "\n", "\r")
+    val startsInline = lineBreaks.none { separator -> completionText.startsWith(separator) }
+
+    if (startsInline) {
+      val renderer =
+          CodyAutocompleteSingleLineRenderer(
+              completionText.lines().first(), items, editor, AutocompleteRendererType.INLINE)
+      inlayModel.addInlineElement(offset, /* relatesToPrecedingText = */ true, renderer)
+    }
+    val lines = completionText.lines()
+    if (lines.size > 1) {
+      val text =
+          (if (startsInline) lines.drop(1) else lines).dropWhile { it.isBlank() }.joinToString("\n")
+      val renderer = CodyAutocompleteBlockElementRenderer(text, items, editor)
       inlayModel.addBlockElement(
-          offset,
-          true,
-          false,
-          Int.MAX_VALUE,
-          CodyAutocompleteBlockElementRenderer(multilineInsertText, items, editor))
+          /* offset = */ offset,
+          /* relatesToPrecedingText = */ true,
+          /* showAbove = */ false,
+          /* priority = */ Int.MAX_VALUE,
+          /* renderer = */ renderer)
     }
   }
 

src/main/kotlin/com/sourcegraph/utils/CodyFormatter.kt

@@ -3,9 +3,9 @@ package com.sourcegraph.utils
 import com.intellij.openapi.editor.Document
 import com.intellij.openapi.fileEditor.FileDocumentManager
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.TextRange
 import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.codeStyle.CodeStyleManager
-import com.intellij.refactoring.suggested.endOffset
 
 class CodyFormatter {
   companion object {
@@ -14,34 +14,37 @@ class CodyFormatter {
      * the document.
      */
     fun formatStringBasedOnDocument(
-        originalText: String,
+        completionText: String,
         project: Project,
         document: Document,
-        offset: Int
+        range: TextRange,
+        cursor: Int
     ): String {
 
-      val appendedString =
-          document.text.substring(0, offset) + originalText + document.text.substring(offset)
+      val beforeCompletion = document.text.substring(0, range.startOffset)
+      val afterCompletion = document.text.substring(range.endOffset)
+      val appendedString = beforeCompletion + completionText + afterCompletion
 
-      val file = FileDocumentManager.getInstance().getFile(document) ?: return originalText
+      val file = FileDocumentManager.getInstance().getFile(document) ?: return completionText
       val psiFile =
           PsiFileFactory.getInstance(project)
               .createFileFromText("TEMP", file.fileType, appendedString)
 
       val codeStyleManager = CodeStyleManager.getInstance(project)
+      codeStyleManager.reformatText(psiFile, cursor, range.startOffset + completionText.length)
 
-      var i = offset
-      var startRefactoringPosition = offset
-      while ((document.text.elementAt(i - 1) == ' ' ||
-          document.text.elementAt(i - 1) == '\n' ||
-          document.text.elementAt(i - 1) == '\t') && i > 0) {
-        startRefactoringPosition = i
-        i--
-      }
-      var endOffset = offset + psiFile.endOffset - document.textLength
-      codeStyleManager.reformatText(psiFile, startRefactoringPosition, endOffset)
-      endOffset = offset + psiFile.endOffset - document.textLength
-      return psiFile.text.substring(startRefactoringPosition, endOffset)
+      // Fix for the IJ formatting bug which removes spaces even before the given formatting range.
+      val existingStart = appendedString.substring(0, cursor)
+      val formattedStart = psiFile.text.substring(0, cursor)
+      val startOfDiff = existingStart.zip(formattedStart).indexOfFirst { (e, f) -> e != f }
+
+      val formattedText =
+          if (startOfDiff != -1) {
+            val addition = formattedStart.substring(startOfDiff)
+            existingStart + addition + psiFile.text.substring(cursor)
+          } else psiFile.text
+      return formattedText.substring(
+          range.startOffset, cursor + formattedText.length - document.textLength)
     }
   }
 }

src/test/kotlin/utils/CodyFormatterTest.kt

@@ -0,0 +1,85 @@
+package utils
+
+import com.intellij.openapi.util.TextRange
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+import com.sourcegraph.utils.CodyFormatter
+import junit.framework.TestCase
+
+class CodyFormatterTest : BasePlatformTestCase() {
+  private val argsString = "String[] args"
+  private val testFileContent =
+      """|
+         |public class HelloWorld {
+         |    public static void main(${argsString}) {
+         |        System.out.println("Hello World!");
+         |        // MAIN
+         |    }
+         |    // CLASS
+         |}"""
+          .trimMargin()
+
+  private var argListOffset = testFileContent.indexOf(argsString)
+  private var insideMainOffset = testFileContent.indexOf("// MAIN")
+  private var insideClassOffset = testFileContent.indexOf("// CLASS")
+
+  private fun formatText(
+      toFormat: String,
+      offset: Int,
+      range: TextRange? = null,
+      fileContent: String = testFileContent
+  ): String {
+    val psiFile = myFixture.addFileToProject("CodyFormatterTest.java", fileContent)
+    return CodyFormatter.formatStringBasedOnDocument(
+        toFormat,
+        myFixture.project,
+        psiFile.viewProvider.document,
+        range ?: TextRange(offset, offset),
+        offset)
+  }
+
+  fun `test single line formatting`() {
+    TestCase.assertEquals("int x = 2;", formatText("int   x =   2;", insideMainOffset))
+  }
+
+  fun `test single line formatting with overlapping range`() {
+    val range = TextRange(argListOffset, argListOffset + argsString.length)
+    // 'String[]   args' is existing text in the editor, so we do not want to reformat it, but we
+    // want to format the rest
+    TestCase.assertEquals(
+        "String[]   args, int n", formatText("String[]   args,   int   n", range.endOffset))
+  }
+
+  fun `test single line formatting to multiline`() {
+    TestCase.assertEquals(
+        """|
+           |    public static int fib(int n) {
+           |        if (n <= 1) {
+           |            return n;
+           |        }
+           |        return fib(n - 1) + fib(n - 2);
+           |    }"""
+            .trimMargin(),
+        formatText(
+            "public static int fib(int n) { if (n <= 1) { return n; } return fib(n-1) + fib(n-2);  }",
+            insideClassOffset))
+  }
+
+  fun `test fix for IJ formatter bug`() {
+    val existingLine = "    public static void test()    "
+    val completion = "$existingLine  { }"
+    val testFileContent =
+        """|public class HelloWorld {
+           |$existingLine
+           |}"""
+            .trimMargin()
+
+    val rangeStart = testFileContent.indexOf(existingLine)
+    val rangeEnd = rangeStart + existingLine.length
+
+    TestCase.assertEquals(
+        """|    public static void test()    {
+           |    }"""
+            .trimMargin(),
+        formatText(completion, rangeEnd, TextRange(rangeStart, rangeEnd), testFileContent))
+  }
+}