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,15 @@ 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.
+    // 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 +354,31 @@ 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))
-    }
+    val deltas = patch.deltas.sortedBy { it.original.position }
+    val rawCompletionText = deltas.joinToString("") { it.revised.lines.joinToString("") }
+
+    val completionText =
+        if (project == null ||
+            System.getProperty("cody.autocomplete.enableFormatting") == "false") {
+          rawCompletionText
+        } else {
+          CodyFormatter.formatStringBasedOnDocument(
+              rawCompletionText, project, editor.document, offset)
+        }
 
-    // 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()) {
+    if (completionText.lines().count() > 1) {
       inlayModel.addBlockElement(
           offset,
           true,
           false,
           Int.MAX_VALUE,
-          CodyAutocompleteBlockElementRenderer(multilineInsertText, items, editor))
+          CodyAutocompleteBlockElementRenderer(completionText, items, editor))
+    } else {
+      inlayModel.addInlineElement(
+          offset,
+          true,
+          CodyAutocompleteSingleLineRenderer(
+              completionText, items, editor, AutocompleteRendererType.INLINE))
     }
   }
 

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

@@ -28,20 +28,11 @@ class CodyFormatter {
           PsiFileFactory.getInstance(project)
               .createFileFromText("TEMP", file.fileType, appendedString)
 
+      fun endOffset() = offset + psiFile.endOffset - document.textLength
       val codeStyleManager = CodeStyleManager.getInstance(project)
+      codeStyleManager.reformatText(psiFile, offset + 1, endOffset())
 
-      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)
+      return psiFile.text.substring(offset, endOffset())
     }
   }
 }

src/test/kotlin/utils/CodyFormatterTest.kt

@@ -0,0 +1,49 @@
+package utils
+
+import com.intellij.ide.highlighter.JavaFileType
+import com.intellij.psi.PsiFileFactory
+import com.intellij.testFramework.fixtures.BasePlatformTestCase
+import com.sourcegraph.utils.CodyFormatter
+import junit.framework.TestCase
+
+class CodyFormatterTest : BasePlatformTestCase() {
+  private val testFileContent =
+      """|
+         |public class HelloWorld {
+         |    public static void main(String[] args) {
+         |        System.out.println("Hello World!");
+         |        // MAIN
+         |    }
+         |    // CLASS
+         |}"""
+          .trimMargin()
+
+  private var insideMainOffset = testFileContent.indexOf("// MAIN")
+  private var insideClassOffset = testFileContent.indexOf("// CLASS")
+
+  private fun formatText(text: String, offset: Int): String {
+    val psiFactory = PsiFileFactory.getInstance(project)
+    val psiFile =
+        psiFactory.createFileFromText("FORMATTING_TEST", JavaFileType.INSTANCE, testFileContent)
+    return CodyFormatter.formatStringBasedOnDocument(
+        text, myFixture.project, psiFile.viewProvider.document, offset)
+  }
+
+  fun `test single line formatting`() {
+    TestCase.assertEquals("int x = 2;", formatText("int   x =   2;", insideMainOffset))
+  }
+
+  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))
+  }
+}