Reuse Parser.parse results (Kotlin#622)

Summary:
Cuts down formatting time on coroutines from 4.7 to ~2.9-3 seconds.
All explanations are in commit messages (note: they should be reviewed separately, the first one addresses a non-functional but necessary change)

Partially addresses Kotlin#619
Fixes Kotlin#552

Pull Request resolved: Kotlin#622

Reviewed By: cortinico

Differential Revision: D108656440

Pulled By: hick209

fbshipit-source-id: b57200ad68139636b1a12d9680dbd5dd0be56195

Author: qwwdfsad

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/).
 ### Changed
 
 * Reduced overall number of allocations to improve formatting performance (~6-7%) (https://github.com/facebook/ktfmt/pull/620)
+* Reuse results of `Parser.parse` (https://github.com/facebook/ktfmt/pull/622)
 
 ## [0.63]
 

core/api/ktfmt.api

@@ -296,7 +296,6 @@ public final class com/facebook/ktfmt/format/KotlinToken : com/google/googlejava
 
 public final class com/facebook/ktfmt/format/MultilineStringFormatter {
 	public fun <init> (I)V
-	public final fun format (Ljava/lang/String;)Ljava/lang/String;
 	public final fun getContinuationIndentSize ()I
 }
 
@@ -325,8 +324,6 @@ public final class com/facebook/ktfmt/format/PsiUtilsKt {
 
 public final class com/facebook/ktfmt/format/RedundantElementManager {
 	public static final field INSTANCE Lcom/facebook/ktfmt/format/RedundantElementManager;
-	public final fun addRedundantElements (Ljava/lang/String;Lcom/facebook/ktfmt/format/FormattingOptions;)Ljava/lang/String;
-	public final fun dropRedundantElements (Ljava/lang/String;Lcom/facebook/ktfmt/format/FormattingOptions;)Ljava/lang/String;
 }
 
 public final class com/facebook/ktfmt/format/Tokenizer : org/jetbrains/kotlin/psi/KtTreeVisitorVoid {

core/src/main/java/com/facebook/ktfmt/format/Formatter.kt

@@ -36,6 +36,7 @@ import org.jetbrains.kotlin.com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiElementVisitor
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtImportDirective
 import org.jetbrains.kotlin.psi.psiUtil.endOffset
 import org.jetbrains.kotlin.psi.psiUtil.startOffset
@@ -98,21 +99,21 @@ object Formatter {
         }
     checkEscapeSequences(kotlinCode)
 
-    return kotlinCode
-        .let { convertLineSeparators(it) }
-        .let { sortedAndDistinctImports(it) }
-        .let { dropRedundantElements(it, options) }
-        .let { addRedundantElements(it, options) }
-        .let { prettyPrint(it, options, lineSeparator = "\n") }
-        .let { addRedundantElements(it, options) }
-        .let { MultilineStringFormatter(options.continuationIndent).format(it) }
+    return FormatterContext(convertLineSeparators(kotlinCode))
+        .transform { sortedAndDistinctImports(it) }
+        .transform { dropRedundantElements(it, options) }
+        .transform { addRedundantElements(it, options) }
+        .transform { prettyPrint(it, options, lineSeparator = "\n") }
+        .transform { addRedundantElements(it, options) }
+        .transform { MultilineStringFormatter(options.continuationIndent).format(it) }
+        .code
         .let { convertLineSeparators(it, checkNotNull(Newlines.guessLineSeparator(kotlinCode))) }
         .let { if (shebang.isEmpty()) it else shebang + "\n" + it }
   }
 
   /** prettyPrint reflows 'code' using google-java-format's engine. */
-  private fun prettyPrint(code: String, options: FormattingOptions, lineSeparator: String): String {
-    val file = Parser.parse(code)
+  private fun prettyPrint(file: KtFile, options: FormattingOptions, lineSeparator: String): String {
+    val code = file.text
     val kotlinInput = KotlinInput(code, file)
     val javaOutput =
         JavaOutput(lineSeparator, kotlinInput, KDocCommentsHelper(lineSeparator, options.maxWidth))
@@ -157,8 +158,8 @@ object Formatter {
     }
   }
 
-  private fun sortedAndDistinctImports(code: String): String {
-    val file = Parser.parse(code)
+  private fun sortedAndDistinctImports(file: KtFile): String {
+    val code = file.text
 
     val importList = file.importList ?: return code
     if (importList.imports.isEmpty()) {
@@ -190,10 +191,20 @@ object Formatter {
     val sortedImports = importList.imports.sortedBy(::canonicalText).distinctBy(::canonicalText)
     val importsWithComments = commentList + sortedImports
 
+    val body = importsWithComments.joinToString(separator = "\n") { imprt -> imprt.text }
+    /*
+     * Kludge: idempotent formatting.
+     * This step optimizes the following goal -- producing **identical** code for already formatted
+     * code, as it's important for PSI-reuse.
+     * There is exactly one case where this step should add trailing newline -- when an inline
+     * comment follows the last import statement. We check for that (note it gives false positives for `/* // */`
+     * which is acceptable -- later prettyPrint step will fix that) and avoid extra-append when it is redundant.
+     */
+    val needsTerminator = body.lastIndexOf('\n').let { it >= 0 && body.indexOf("//", it + 1) >= 0 }
     return code.replaceRange(
         importList.startOffset,
         importList.endOffset,
-        importsWithComments.joinToString(separator = "\n") { imprt -> imprt.text } + "\n",
+        if (needsTerminator) body + "\n" else body,
     )
   }
 }

core/src/main/java/com/facebook/ktfmt/format/FormatterContext.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.ktfmt.format
+
+import org.jetbrains.kotlin.psi.KtFile
+
+internal class FormatterContext(@JvmField val code: String) {
+
+  private val ktFile: KtFile by lazy { Parser.parse(code) }
+
+  inline fun transform(block: (KtFile) -> String): FormatterContext {
+    val newCode = block(ktFile)
+    return if (newCode == code) this else FormatterContext(newCode)
+  }
+}

core/src/main/java/com/facebook/ktfmt/format/MultilineStringFormatter.kt

@@ -18,6 +18,7 @@ package com.facebook.ktfmt.format
 
 import com.google.common.annotations.VisibleForTesting
 import org.jetbrains.kotlin.com.intellij.psi.PsiComment
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtQualifiedExpression
 import org.jetbrains.kotlin.psi.KtStringTemplateExpression
 import org.jetbrains.kotlin.psi.KtTreeVisitorVoid
@@ -31,10 +32,11 @@ private const val TQ = "\"\"\""
  * imports.
  */
 class MultilineStringFormatter(val continuationIndentSize: Int) {
-  fun format(code: String): String {
+  internal fun format(file: KtFile): String {
+    val code = file.text
     val result = StringBuilder(code)
     val multilineStringList =
-        getMultilineTrimmedStringList(code)
+        getMultilineTrimmedStringList(file)
             .sortedByDescending(MultilineTrimmedString::openStringOffset)
     for (multilineString in multilineStringList) {
       if (multilineString.stringLineCount < 2) {
@@ -104,8 +106,8 @@ class MultilineStringFormatter(val continuationIndentSize: Int) {
   }
 
   @VisibleForTesting
-  internal fun getMultilineTrimmedStringList(code: String): List<MultilineTrimmedString> {
-    val file = Parser.parse(code)
+  internal fun getMultilineTrimmedStringList(file: KtFile): List<MultilineTrimmedString> {
+    val code = file.text
     val strings = mutableListOf<MultilineTrimmedString>()
     file.accept(
         object : KtTreeVisitorVoid() {

core/src/main/java/com/facebook/ktfmt/format/RedundantElementManager.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
 import org.jetbrains.kotlin.kdoc.psi.impl.KDocImpl
 import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.psi.KtImportList
 import org.jetbrains.kotlin.psi.KtPackageDirective
 import org.jetbrains.kotlin.psi.KtReferenceExpression
@@ -33,8 +34,8 @@ import org.jetbrains.kotlin.psi.psiUtil.startOffset
  */
 object RedundantElementManager {
   /** Remove extra semicolons and unused imports, if enabled in the [options] */
-  fun dropRedundantElements(code: String, options: FormattingOptions): String {
-    val file = Parser.parse(code)
+  internal fun dropRedundantElements(file: KtFile, options: FormattingOptions): String {
+    val code = file.text
     val redundantImportDetector = RedundantImportDetector(enabled = options.removeUnusedImports)
     val redundantSemicolonDetector = RedundantSemicolonDetector()
     val trailingCommaDetector = TrailingCommas.Detector()
@@ -92,12 +93,12 @@ object RedundantElementManager {
     return result.toString()
   }
 
-  fun addRedundantElements(code: String, options: FormattingOptions): String {
+  internal fun addRedundantElements(file: KtFile, options: FormattingOptions): String {
     if (!options.manageTrailingCommas) {
-      return code
+      return file.text
     }
 
-    val file = Parser.parse(code)
+    val code = file.text
     val trailingCommaSuggestor = TrailingCommas.Suggestor()
 
     file.accept(

core/src/test/java/com/facebook/ktfmt/format/FormatterTest.kt

@@ -1219,6 +1219,48 @@ class FormatterTest {
     assertThatFormatting(code).isEqualTo(expected)
   }
 
+  @Test
+  fun `imports with trailing block comment and expression`() {
+    val code =
+        """
+        |import com.example.zab /* heya */
+        |import com.example.foo ; val x = Sample(foo, zab)
+        |"""
+            .trimMargin()
+
+    val expected =
+        """
+        |import com.example.foo
+        |import com.example.zab /* heya */
+        |
+        |val x = Sample(foo, zab)
+        |"""
+            .trimMargin()
+
+    assertThatFormatting(code).isEqualTo(expected)
+  }
+
+  @Test
+  fun `imports with trailing block comment and expression with false positive double slash`() {
+    val code =
+        """
+        |import com.example.zab /* // */
+        |import com.example.foo ; val x = Sample(foo, zab)
+        |"""
+            .trimMargin()
+
+    val expected =
+        """
+        |import com.example.foo
+        |import com.example.zab /* // */
+        |
+        |val x = Sample(foo, zab)
+        |"""
+            .trimMargin()
+
+    assertThatFormatting(code).isEqualTo(expected)
+  }
+
   @Test
   fun `backticks are ignored in import sort order`() =
       assertFormatted(

core/src/test/java/com/facebook/ktfmt/format/MultilineStringFormatterTest.kt

@@ -303,7 +303,9 @@ class MultilineStringFormatterTest {
       @Language("kts") code: String,
       continuationIndent: Int = 4,
   ): MultilineTrimmedString {
-    val strings = MultilineStringFormatter(continuationIndent).getMultilineTrimmedStringList(code)
+    val strings =
+        MultilineStringFormatter(continuationIndent)
+            .getMultilineTrimmedStringList(Parser.parse(code))
     assertThat(strings.size).isEqualTo(1)
     return strings.first()
   }