Avoid redundant allocations (Kotlin#620)

Summary:
* Optimizes work with `PsiElement`: no direct `text` access where appropriate as it's relatively expensive
* Some early returns to avoid String copying when it is not necessary
* Around 6-7% performance improvements on my machine on sequential `kotlinx.coroutines` runs (5 -> 4.65 sec in sequantial mode)

Partially addresses Kotlin#619

Pull Request resolved: Kotlin#620

Reviewed By: cortinico

Differential Revision: D107013190

Pulled By: hick209

fbshipit-source-id: 8a420697a5d5bbe708e46d016c0042b9d2fd3a71

Author: qwwdfsad

CHANGELOG.md

@@ -14,6 +14,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)
 
 ## [0.63]
 

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

@@ -130,8 +130,11 @@ class KotlinInput(private val text: String, file: KtFile) : Input() {
         )
   }
 
-  private fun makePositionToColumnMap(toks: List<KotlinTok>) =
-      ImmutableMap.copyOf(toks.map { it.position to it.column }.toMap())
+  private fun makePositionToColumnMap(toks: List<KotlinTok>): ImmutableMap<Int, Int> {
+    val builder = ImmutableMap.builderWithExpectedSize<Int, Int>(toks.size)
+    toks.forEach { builder.put(it.position, it.column) }
+    return builder.build()
+  }
 
   private fun buildToks(file: KtFile, fileText: String): ImmutableList<KotlinTok> {
     val tokenizer = Tokenizer(fileText, file)

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

@@ -1619,7 +1619,7 @@ class KotlinInputAstVisitor(
   private fun hasSourceNewlineInLambdaBody(lambdaExpression: KtLambdaExpression): Boolean {
     val functionLiteral = lambdaExpression.functionLiteral
     for (child in functionLiteral.node.children()) {
-      if (child.psi is PsiWhiteSpace && child.text.contains('\n')) return true
+      if (child.psi is PsiWhiteSpace && child.textContains('\n')) return true
     }
     return false
   }

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

@@ -71,11 +71,12 @@ object RedundantElementManager {
         }
     )
 
-    val result = StringBuilder(code)
     val elementsToRemove =
         redundantSemicolonDetector.getRedundantSemicolonElements() +
             redundantImportDetector.getRedundantImportElements() +
             trailingCommaDetector.getTrailingCommaElements()
+    if (elementsToRemove.isEmpty()) return code
+    val result = StringBuilder(code)
 
     for (element in elementsToRemove.sortedByDescending(PsiElement::endOffset)) {
       // Don't insert extra newlines when the semicolon is already a line terminator
@@ -108,8 +109,9 @@ object RedundantElementManager {
         }
     )
 
-    val result = StringBuilder(code)
     val suggestionElements = trailingCommaSuggestor.getTrailingCommaSuggestions()
+    if (suggestionElements.isEmpty()) return code
+    val result = StringBuilder(code)
 
     for (element in suggestionElements.sortedByDescending(PsiElement::endOffset)) {
       result.insert(element.endOffset, ',')
@@ -120,6 +122,6 @@ object RedundantElementManager {
 
   private fun PsiElement?.containsNewline(): Boolean {
     if (this !is PsiWhiteSpace) return false
-    return this.text.contains('\n')
+    return this.textContains('\n')
   }
 }

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

@@ -52,7 +52,7 @@ internal class RedundantSemicolonDetector {
 
   /** returns **true** if this element was an extra comma, **false** otherwise. */
   private fun isExtraSemicolon(element: PsiElement): Boolean {
-    if (element.text != ";") {
+    if (element !is LeafPsiElement || element.elementType != KtTokens.SEMICOLON) {
       return false
     }
 

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

@@ -61,13 +61,16 @@ class Tokenizer(private val fileText: String, val file: KtFile) : KtTreeVisitorV
   }
 
   override fun visitElement(element: PsiElement) {
-    val startIndex = element.startOffset
-    val endIndex = element.endOffset
-    val elementText = element.text
-    val originalText = fileText.substring(startIndex, endIndex)
+    // Do not materialize text/text ranges when it's non needed -- e.g. for KtFile, composite
+    // PsiElement(...) etc.
+    val elementText by lazy(LazyThreadSafetyMode.NONE) { element.text }
+    val originalText by
+        lazy(LazyThreadSafetyMode.NONE) {
+          fileText.substring(element.startOffset, element.endOffset)
+        }
     when (element) {
       is PsiComment -> {
-        if (element.text.startsWith("/*") && !element.text.endsWith("*/")) {
+        if (elementText.startsWith("/*") && !elementText.endsWith("*/")) {
           throw ParseError(
               "Unclosed comment",
               StringUtil.offsetToLineColumn(fileText, element.startOffset),
@@ -87,7 +90,7 @@ class Tokenizer(private val fileText: String, val file: KtFile) : KtTreeVisitorV
                 index = index,
                 originalText = originalText,
                 text = elementText,
-                position = startIndex,
+                position = element.startOffset,
                 column = 0,
                 isToken = treatAsToken,
                 kind = KtTokens.EOF,
@@ -105,7 +108,7 @@ class Tokenizer(private val fileText: String, val file: KtFile) : KtTreeVisitorV
                         originalText,
                     ),
                 text = elementText,
-                position = startIndex,
+                position = element.startOffset,
                 column = 0,
                 isToken = true,
                 kind = KtTokens.EOF,
@@ -124,11 +127,11 @@ class Tokenizer(private val fileText: String, val file: KtFile) : KtTreeVisitorV
                     index = -1,
                     originalText =
                         fileText.substring(
-                            startIndex + matcher.start(),
-                            startIndex + matcher.end(),
+                            element.startOffset + matcher.start(),
+                            element.startOffset + matcher.end(),
                         ),
                     text = text,
-                    position = startIndex + matcher.start(),
+                    position = element.startOffset + matcher.start(),
                     column = 0,
                     isToken = false,
                     kind = KtTokens.EOF,
@@ -141,7 +144,7 @@ class Tokenizer(private val fileText: String, val file: KtFile) : KtTreeVisitorV
                   index = index,
                   originalText = originalText,
                   text = elementText,
-                  position = startIndex,
+                  position = element.startOffset,
                   column = 0,
                   isToken = true,
                   kind = KtTokens.EOF,

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

@@ -19,6 +19,8 @@ package com.facebook.ktfmt.format
 import org.jetbrains.kotlin.com.intellij.psi.PsiComment
 import org.jetbrains.kotlin.com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.com.intellij.psi.PsiWhiteSpace
+import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafPsiElement
+import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.KtClassBody
 import org.jetbrains.kotlin.psi.KtCollectionLiteralExpression
 import org.jetbrains.kotlin.psi.KtElement
@@ -47,10 +49,9 @@ object TrailingCommas {
     }
 
     private fun isTrailingComma(element: PsiElement): Boolean {
-      if (element.text != ",") {
+      if (element !is LeafPsiElement || element.elementType != KtTokens.COMMA) {
         return false
       }
-
       return extractManagedList(element.parent)?.trailingComma == element
     }
   }
@@ -77,10 +78,6 @@ object TrailingCommas {
      * ```
      */
     fun takeElement(element: KtElement) {
-      if (!element.text.contains("\n")) {
-        return // Only suggest trailing commas where there is already a line break
-      }
-
       when (element) {
         is KtEnumEntry, // Only suggest on the KtClassBody container
         is KtWhenEntry -> return
@@ -99,6 +96,9 @@ object TrailingCommas {
       }
 
       val list = extractManagedList(element) ?: return
+      if (!element.textContains('\n')) {
+        return // Only suggest trailing commas where there is already a line break
+      }
       if (list.items.size <= 1) {
         return // Never insert commas to single-element lists
       }

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

@@ -22,6 +22,7 @@ import java.util.regex.Pattern.MULTILINE
 object WhitespaceTombstones {
   /** See [replaceTrailingWhitespaceWithTombstone]. */
   const val SPACE_TOMBSTONE = '\u0003'
+  private val TRAILING_WHITESPACE_PATTERN = Pattern.compile(" ($)", MULTILINE)
 
   fun String.indexOfWhitespaceTombstone(): Int = this.indexOf(SPACE_TOMBSTONE)
 
@@ -32,7 +33,7 @@ object WhitespaceTombstones {
    * replace it back to a space.
    */
   fun replaceTrailingWhitespaceWithTombstone(s: String): String {
-    return Pattern.compile(" ($)", MULTILINE).matcher(s).replaceAll("$SPACE_TOMBSTONE$1")
+    return TRAILING_WHITESPACE_PATTERN.matcher(s).replaceAll("$SPACE_TOMBSTONE$1")
   }
 
   /** See [replaceTrailingWhitespaceWithTombstone]. */