KtBlockExpression as lazy reparseable node (KT-13841)

Changed KtBlockExpression to support LazyReparseable behaviour
Added LazyKtBlockExpressionTest
Added new BlockWrapper delegations

compiler/psi/src/org/jetbrains/kotlin/BlockExpressionElementType.kt

@@ -0,0 +1,107 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin
+
+import com.intellij.lang.ASTNode
+import com.intellij.lang.Language
+import com.intellij.lang.PsiBuilderFactory
+import com.intellij.openapi.project.Project
+import com.intellij.psi.tree.ICompositeElementType
+import com.intellij.psi.tree.IErrorCounterReparseableElementType
+import com.intellij.psi.tree.TokenSet
+import org.jetbrains.kotlin.KtNodeTypes.BLOCK_CODE_FRAGMENT
+import org.jetbrains.kotlin.KtNodeTypes.FUNCTION_LITERAL
+import org.jetbrains.kotlin.KtNodeTypes.SCRIPT
+import org.jetbrains.kotlin.idea.KotlinLanguage
+import org.jetbrains.kotlin.lexer.KotlinLexer
+import org.jetbrains.kotlin.lexer.KtTokens
+import org.jetbrains.kotlin.parsing.KotlinParser
+import org.jetbrains.kotlin.psi.KtBlockExpression
+
+class BlockExpressionElementType : IErrorCounterReparseableElementType("BLOCK", KotlinLanguage.INSTANCE), ICompositeElementType {
+
+    override fun createCompositeNode() = KtBlockExpression(null)
+
+    override fun createNode(text: CharSequence?) = KtBlockExpression(text)
+
+    override fun isParsable(parent: ASTNode?, buffer: CharSequence, fileLanguage: Language, project: Project) =
+        fileLanguage == KotlinLanguage.INSTANCE &&
+                BlockExpressionElementType.isAllowedParentNode(parent) &&
+                BlockExpressionElementType.isReparseableBlock(buffer) &&
+                super.isParsable(buffer, fileLanguage, project)
+
+    override fun getErrorsCount(seq: CharSequence, fileLanguage: Language, project: Project) =
+        ElementTypeUtils.getKotlinBlockImbalanceCount(seq)
+
+    override fun parseContents(chameleon: ASTNode): ASTNode {
+        val project = chameleon.psi.project
+        val builder = PsiBuilderFactory.getInstance().createBuilder(
+            project, chameleon, null, KotlinLanguage.INSTANCE, chameleon.chars
+        )
+
+        return KotlinParser.parseBlockExpression(builder).firstChildNode
+    }
+
+    companion object {
+
+        private fun isAllowedParentNode(node: ASTNode?) =
+            node != null &&
+                    SCRIPT != node.elementType &&
+                    FUNCTION_LITERAL != node.elementType &&
+                    BLOCK_CODE_FRAGMENT != node.elementType
+
+        /**
+         * Check if this text is block but not a lambda, please refer to parsing rules!
+        @see [org.jetbrains.kotlin.parsing.KotlinExpressionParsing.parseFunctionLiteral]
+         */
+        fun isReparseableBlock(blockText: CharSequence): Boolean {
+
+            fun advanceWhitespacesCheckIsEndOrArrow(lexer: KotlinLexer): Boolean {
+                lexer.advance()
+                while (lexer.tokenType != null && lexer.tokenType != KtTokens.EOF) {
+                    if (lexer.tokenType == KtTokens.ARROW) return true
+                    if (lexer.tokenType != KtTokens.WHITE_SPACE) return false
+                    lexer.advance()
+                }
+                return true
+            }
+
+            val lexer = KotlinLexer()
+            lexer.start(blockText)
+
+            // Try to parse a simple name list followed by an ARROW
+            //   {a -> ...}
+            //   {a, b -> ...}
+            //   {(a, b) -> ... }
+            if (lexer.tokenType != KtTokens.LBRACE) return false
+
+            if (advanceWhitespacesCheckIsEndOrArrow(lexer)) return false
+
+            if (lexer.tokenType != KtTokens.COLON &&
+                lexer.tokenType != KtTokens.IDENTIFIER &&
+                lexer.tokenType != KtTokens.LPAR
+            ) return true
+
+            val searchForRPAR = lexer.tokenType == KtTokens.LPAR
+
+            if (advanceWhitespacesCheckIsEndOrArrow(lexer)) return false
+
+            val preferParamsToExpressions = lexer.tokenType == KtTokens.COMMA || lexer.tokenType == KtTokens.COLON
+
+            while (true) {
+
+                if (lexer.tokenType == KtTokens.LBRACE) return true
+                if (lexer.tokenType == KtTokens.RBRACE) return !preferParamsToExpressions
+
+                if (searchForRPAR && lexer.tokenType == KtTokens.RPAR) {
+                    return !advanceWhitespacesCheckIsEndOrArrow(lexer)
+                }
+
+                if (advanceWhitespacesCheckIsEndOrArrow(lexer)) return false
+            }
+        }
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/ElementTypeUtils.kt

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin
+
+import com.intellij.psi.tree.IErrorCounterReparseableElementType
+import org.jetbrains.kotlin.lexer.KotlinLexer
+import org.jetbrains.kotlin.lexer.KtTokens
+
+object ElementTypeUtils {
+    @JvmStatic
+    fun getKotlinBlockImbalanceCount(seq: CharSequence): Int {
+        val lexer = KotlinLexer()
+
+        lexer.start(seq)
+        if (lexer.tokenType !== KtTokens.LBRACE) return IErrorCounterReparseableElementType.FATAL_ERROR
+        lexer.advance()
+        var balance = 1
+        while (lexer.tokenType != KtTokens.EOF) {
+            val type = lexer.tokenType ?: break
+            if (balance == 0) {
+                return IErrorCounterReparseableElementType.FATAL_ERROR
+            }
+            if (type === KtTokens.LBRACE) {
+                balance++
+            } else if (type === KtTokens.RBRACE) {
+                balance--
+            }
+            lexer.advance()
+        }
+        return balance
+    }
+}

compiler/psi/src/org/jetbrains/kotlin/KtNodeTypes.java

@@ -112,9 +112,10 @@ public interface KtNodeTypes {
     IElementType DO_WHILE                  = new KtNodeType("DO_WHILE", KtDoWhileExpression.class);
     IElementType LOOP_RANGE                = new KtNodeType("LOOP_RANGE", KtContainerNode.class);
     IElementType BODY                      = new KtNodeType("BODY", KtContainerNodeForControlStructureBody.class);
-    IElementType BLOCK                     = new KtNodeType("BLOCK", KtBlockExpression.class);
 
-    IElementType LAMBDA_EXPRESSION       = new LambdaExpressionElementType();
+    IElementType BLOCK                     = new BlockExpressionElementType();
+
+    IElementType LAMBDA_EXPRESSION         = new LambdaExpressionElementType();
 
     IElementType FUNCTION_LITERAL          = new KtNodeType("FUNCTION_LITERAL", KtFunctionLiteral.class);
     IElementType ANNOTATED_EXPRESSION      = new KtNodeType("ANNOTATED_EXPRESSION", KtAnnotatedExpression.class);

compiler/psi/src/org/jetbrains/kotlin/LambdaExpressionElementType.java

@@ -147,27 +147,7 @@ class LambdaExpressionElementType extends IErrorCounterReparseableElementType {
     }
 
     @Override
-    public int getErrorsCount(CharSequence seq, Language fileLanguage, Project project) {
+    public int getErrorsCount(CharSequence seq, Language fileLanguage, Project project){
-        Lexer lexer = new KotlinLexer();
+        return ElementTypeUtils.getKotlinBlockImbalanceCount(seq);
-
-        lexer.start(seq);
-        if (lexer.getTokenType() != KtTokens.LBRACE) return IErrorCounterReparseableElementType.FATAL_ERROR;
-        lexer.advance();
-        int balance = 1;
-        while (true) {
-            IElementType type = lexer.getTokenType();
-            if (type == null) break;
-            if (balance == 0) {
-                return IErrorCounterReparseableElementType.FATAL_ERROR;
-            }
-            if (type == KtTokens.LBRACE) {
-                balance++;
-            }
-            else if (type == KtTokens.RBRACE) {
-                balance--;
-            }
-            lexer.advance();
-        }
-        return balance;
     }
 }

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinExpressionParsing.java

@@ -17,7 +17,10 @@
 package org.jetbrains.kotlin.parsing;
 
 import com.google.common.collect.ImmutableMap;
+import com.intellij.lang.ASTNode;
+import com.intellij.lang.Language;
 import com.intellij.lang.PsiBuilder;
+import com.intellij.openapi.project.Project;
 import com.intellij.psi.tree.IElementType;
 import com.intellij.psi.tree.TokenSet;
 import org.jetbrains.annotations.NotNull;
@@ -1093,6 +1096,8 @@ public class KotlinExpressionParsing extends AbstractKotlinParsing {
 
     /**
      * If it has no ->, it's a block, otherwise a function literal
+     *
+     * Please update {@link org.jetbrains.kotlin.BlockExpressionElementType#isParsable(ASTNode, CharSequence, Language, Project)} if any changes occurs!
      */
     public void parseFunctionLiteral(boolean preferBlock, boolean collapse) {
         assert _at(LBRACE);
@@ -1138,7 +1143,7 @@ public class KotlinExpressionParsing extends AbstractKotlinParsing {
         }
 
         if (collapse) {
-            advanceLambdaBlock();
+            myKotlinParsing.advanceBalancedBlock();
             literal.done(FUNCTION_LITERAL);
             literalExpression.collapse(LAMBDA_EXPRESSION);
         }
@@ -1157,24 +1162,6 @@ public class KotlinExpressionParsing extends AbstractKotlinParsing {
         myBuilder.restoreNewlinesState();
     }
 
-    private void advanceLambdaBlock() {
-        int braceCount = 1;
-        while (!eof()) {
-            if (_at(LBRACE)) {
-                braceCount++;
-            }
-            else if (_at(RBRACE)) {
-                braceCount--;
-            }
-
-            advance();
-
-            if (braceCount == 0) {
-                break;
-            }
-        }
-    }
-
     private boolean rollbackOrDropAt(PsiBuilder.Marker rollbackMarker, IElementType dropAt) {
         if (at(dropAt)) {
             advance(); // dropAt

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinParser.java

@@ -80,4 +80,11 @@ public class KotlinParser implements PsiParser {
         ktParsing.parseLambdaExpression();
         return psiBuilder.getTreeBuilt();
     }
+
+    @NotNull
+    public static ASTNode parseBlockExpression(PsiBuilder psiBuilder) {
+        KotlinParsing ktParsing = KotlinParsing.createForTopLevel(new SemanticWhitespaceAwarePsiBuilderImpl(psiBuilder));
+        ktParsing.parseBlockExpression();
+        return psiBuilder.getTreeBuilt();
+    }
 }

compiler/psi/src/org/jetbrains/kotlin/parsing/KotlinParsing.java

@@ -162,6 +162,10 @@ public class KotlinParsing extends AbstractKotlinParsing {
         myExpressionParsing.parseFunctionLiteral(/* preferBlock = */ false, /* collapse = */false);
     }
 
+    void parseBlockExpression() {
+        parseBlock(/* collapse = */ false);
+    }
+
     void parseScript() {
         PsiBuilder.Marker fileMarker = mark();
 
@@ -1725,17 +1729,50 @@ public class KotlinParsing extends AbstractKotlinParsing {
      *   ;
      */
     void parseBlock() {
-        PsiBuilder.Marker block = mark();
+        parseBlock(/*collapse*/ true);
+    }
+
+    private void parseBlock(boolean collapse) {
+
+        PsiBuilder.Marker lazyBlock = mark();
 
         myBuilder.enableNewlines();
+
         expect(LBRACE, "Expecting '{' to open a block");
 
-        myExpressionParsing.parseStatements();
+        if(collapse){
+            advanceBalancedBlock();
+        }else{
+            myExpressionParsing.parseStatements();
+            expect(RBRACE, "Expecting '}'");
+        }
 
-        expect(RBRACE, "Expecting '}'");
         myBuilder.restoreNewlinesState();
 
-        block.done(BLOCK);
+        if(collapse){
+            lazyBlock.collapse(BLOCK);
+        }else{
+            lazyBlock.done(BLOCK);
+        }
+    }
+
+    public void advanceBalancedBlock() {
+
+        int braceCount = 1;
+        while (!eof()) {
+            if (_at(LBRACE)) {
+                braceCount++;
+            }
+            else if (_at(RBRACE)) {
+                braceCount--;
+            }
+
+            advance();
+
+            if (braceCount == 0) {
+                break;
+            }
+        }
     }
 
     /*

compiler/psi/src/org/jetbrains/kotlin/psi/KtBlockExpression.java

@@ -16,21 +16,32 @@
 
 package org.jetbrains.kotlin.psi;
 
-import com.intellij.lang.ASTNode;
+import com.intellij.lang.Language;
 import com.intellij.openapi.util.TextRange;
-import com.intellij.psi.PsiElement;
+import com.intellij.psi.*;
-import com.intellij.psi.PsiModifiableCodeBlock;
+import com.intellij.psi.impl.source.resolve.reference.ReferenceProvidersRegistry;
+import com.intellij.psi.impl.source.tree.CompositeElement;
+import com.intellij.psi.impl.source.tree.LazyParseablePsiElement;
+import com.intellij.psi.util.PsiUtilCore;
+import com.intellij.util.IncorrectOperationException;
 import kotlin.annotations.jvm.ReadOnly;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.jetbrains.kotlin.idea.KotlinLanguage;
 import org.jetbrains.kotlin.lexer.KtTokens;
+import org.jetbrains.kotlin.psi.psiUtil.KtPsiUtilKt;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilsKt;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 
-public class KtBlockExpression extends KtExpressionImpl implements KtStatementExpression, PsiModifiableCodeBlock {
+import static org.jetbrains.kotlin.KtNodeTypes.BLOCK;
-    public KtBlockExpression(@NotNull ASTNode node) {
+
-        super(node);
+public class KtBlockExpression extends LazyParseablePsiElement implements KtElement, KtExpression, KtStatementExpression, PsiModifiableCodeBlock {
+
+    public KtBlockExpression(@Nullable CharSequence text) {
+        super(BLOCK, text);
     }
 
     @Override
@@ -39,11 +50,92 @@ public class KtBlockExpression extends KtExpressionImpl implements KtStatementEx
         return false;
     }
 
+    @NotNull
+    @Override
+    public Language getLanguage() {
+        return KotlinLanguage.INSTANCE;
+    }
+
+    @Override
+    public String toString() {
+        return getNode().getElementType().toString();
+    }
+
+    @NotNull
+    @Override
+    public KtFile getContainingKtFile() {
+        return PsiUtilsKt.getContainingKtFile(this);
+    }
+
+    @Override
+    public <D> void acceptChildren(@NotNull KtVisitor<Void, D> visitor, D data) {
+        KtPsiUtil.visitChildren(this, visitor, data);
+    }
+
     @Override
     public <R, D> R accept(@NotNull KtVisitor<R, D> visitor, D data) {
         return visitor.visitBlockExpression(this, data);
     }
 
+    @Override
+    @SuppressWarnings("unchecked")
+    public final void accept(@NotNull PsiElementVisitor visitor) {
+        if (visitor instanceof KtVisitor) {
+            accept((KtVisitor) visitor, null);
+        }
+        else {
+            visitor.visitElement(this);
+        }
+    }
+
+    @Override
+    public void delete() throws IncorrectOperationException {
+        KtElementUtilsKt.deleteSemicolon(this);
+        super.delete();
+    }
+
+    @Override
+    @SuppressWarnings("deprecation")
+    public PsiReference getReference() {
+        PsiReference[] references = getReferences();
+        if (references.length == 1) return references[0];
+        else return null;
+    }
+
+    @Override
+    @NotNull
+    public PsiElement[] getChildren() {
+        PsiElement psiChild = getFirstChild();
+
+        List<PsiElement> result = null;
+        while (psiChild != null) {
+            if (psiChild.getNode() instanceof CompositeElement) {
+                if(result == null) result = new ArrayList<>();
+                result.add(psiChild);
+            }
+            psiChild = psiChild.getNextSibling();
+        }
+        return result == null ? PsiElement.EMPTY_ARRAY : PsiUtilCore.toPsiElementArray(result);
+    }
+
+    @NotNull
+    @Override
+    public PsiReference[] getReferences() {
+        return ReferenceProvidersRegistry.getReferencesFromProviders(this, PsiReferenceService.Hints.NO_HINTS);
+    }
+
+    @NotNull
+    @Override
+    public KtElement getPsiOrParent() {
+        return this;
+    }
+
+    @Override
+    public PsiElement getParent() {
+        PsiElement substitute = KtPsiUtilKt.getParentSubstitute(this);
+        return substitute != null ? substitute : super.getParent();
+    }
+
     @ReadOnly
     @NotNull
     public List<KtExpression> getStatements() {
@@ -58,11 +150,11 @@ public class KtBlockExpression extends KtExpressionImpl implements KtStatementEx
 
     @Nullable
     public PsiElement getRBrace() {
-        return findChildByType(KtTokens.RBRACE);
+        return findPsiChildByType(KtTokens.RBRACE);
     }
 
     @Nullable
     public PsiElement getLBrace() {
-        return findChildByType(KtTokens.LBRACE);
+        return findPsiChildByType(KtTokens.LBRACE);
     }
 }

compiler/psi/src/org/jetbrains/kotlin/psi/KtLambdaExpression.java

@@ -26,6 +26,7 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.kotlin.KtNodeTypes;
 import org.jetbrains.kotlin.lexer.KtTokens;
+import org.jetbrains.kotlin.psi.psiUtil.PsiUtilsKt;
 
 import java.util.List;
 
@@ -76,13 +77,7 @@ public class KtLambdaExpression extends LazyParseablePsiElement implements KtExp
     @NotNull
     @Override
     public KtFile getContainingKtFile() {
-        PsiFile file = getContainingFile();
+        return PsiUtilsKt.getContainingKtFile(this);
-        if(!(file instanceof KtFile))  {
-            String fileString = (file != null && file.isValid()) ? file.getText() : "";
-            throw new IllegalStateException("KtElement not inside KtFile: " + file + fileString +
-                                            "for element " + this + " of type " + this.getClass() + " node = " + getNode());
-        }
-        return (KtFile) file;
     }
 
     @Override

compiler/psi/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiComment
 import com.intellij.psi.PsiElement
+import com.intellij.psi.PsiFile
 import com.intellij.psi.PsiFileFactory
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.util.LocalTimeCounter
@@ -853,7 +854,8 @@ class KtPsiFactory @JvmOverloads constructor(private val project: Project, val m
     }
 
     private class BlockWrapper(fakeBlockExpression: KtBlockExpression, private val expression: KtExpression) :
-        KtBlockExpression(fakeBlockExpression.node), KtPsiUtil.KtExpressionWrapper {
+        KtBlockExpression(fakeBlockExpression.text), KtPsiUtil.KtExpressionWrapper {
+
         override fun getStatements(): List<KtExpression> {
             return listOf(expression)
         }
@@ -861,5 +863,13 @@ class KtPsiFactory @JvmOverloads constructor(private val project: Project, val m
         override fun getBaseExpression(): KtExpression {
             return expression
         }
+
+        override fun getParent(): PsiElement = expression.parent
+
+        override fun getPsiOrParent(): KtElement = expression.psiOrParent
+
+        override fun getContainingKtFile() = expression.containingKtFile
+
+        override fun getContainingFile(): PsiFile = expression.containingFile
     }
 }

compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/psiUtils.kt

@@ -20,6 +20,7 @@ import com.intellij.injected.editor.VirtualFileWindow
 import com.intellij.lang.ASTNode
 import com.intellij.openapi.util.TextRange
 import com.intellij.psi.*
+import com.intellij.psi.impl.source.tree.LazyParseablePsiElement
 import com.intellij.psi.impl.source.tree.TreeUtil
 import com.intellij.psi.search.PsiSearchScopeUtil
 import com.intellij.psi.search.SearchScope
@@ -380,7 +381,7 @@ fun PsiElement.getElementTextWithContext(): String {
 
     // Find parent for element among file children
     val topLevelElement = PsiTreeUtil.findFirstParent(this, { it.parent is PsiFile })
-            ?: throw AssertionError("For non-file element we should always be able to find parent in file children")
+        ?: throw AssertionError("For non-file element we should always be able to find parent in file children")
 
     val startContextOffset = topLevelElement.startOffset
     val elementContextOffset = textRange.startOffset
@@ -466,3 +467,13 @@ fun ASTNode.closestPsiElement(): PsiElement? {
     }
     return node.psi
 }
+
+fun LazyParseablePsiElement.getContainingKtFile(): KtFile {
+
+    val file = this.containingFile
+
+    if (file is KtFile) return file
+
+    val fileString = if (file != null && file.isValid) file.text else ""
+    throw IllegalStateException("KtElement not inside KtFile: $file with text \"$fileString\" for element $this of type ${this::class.java} node = ${this.node}")
+}

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/CodeFragmentAnalyzer.kt

@@ -108,7 +108,7 @@ class CodeFragmentAnalyzer(
                 }
             }
             is KtSecondaryConstructor -> {
-                val expression = context.bodyExpression ?: context.getDelegationCall().calleeExpression
+                val expression = (context.bodyExpression ?: context.getDelegationCall().calleeExpression) as? KtExpression
                 if (expression != null) {
                     bindingContext = resolutionFactory(expression)
                     scope = bindingContext[BindingContext.LEXICAL_SCOPE, expression]

idea/tests/org/jetbrains/kotlin/idea/editor/LazyElementTypeTest.kt

@@ -1,62 +0,0 @@
-/*
- * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
- * that can be found in the license/LICENSE.txt file.
- */
-
-package org.jetbrains.kotlin.idea.editor
-
-import com.intellij.psi.PsiDocumentManager
-import com.intellij.psi.util.PsiTreeUtil
-import com.intellij.testFramework.EditorTestUtil
-import com.intellij.testFramework.EditorTestUtil.*
-import com.intellij.testFramework.LightPlatformTestCase
-import com.intellij.testFramework.LightProjectDescriptor
-import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCaseBase
-import org.jetbrains.kotlin.psi.KtLambdaExpression
-import org.jetbrains.kotlin.test.JUnit3WithIdeaConfigurationRunner
-import org.junit.Assert
-import org.junit.runner.RunWith
-
-@RunWith(JUnit3WithIdeaConfigurationRunner::class)
-class LazyElementTypeTest : KotlinLightCodeInsightFixtureTestCaseBase() {
-    fun testSplitArrow() = reparse("val t = { a: Int -<caret>> }", ' ')
-    fun testDeleteArrow() = reparse("val t = { a: Int -><caret> }", BACKSPACE_FAKE_CHAR)
-
-    fun testReformatNearArrow() = noReparse("val t = { a: Int<caret>-> }", ' ')
-    fun testChangeAfterArrow() = noReparse("val t = { a: Int -> <caret> }", 'a')
-    fun testDeleteIrrelevantArrow() = noReparse("val t = { a: Int -> (1..3).filter { b -><caret> b > 2 } }", BACKSPACE_FAKE_CHAR)
-    fun testReformatNearLambdaStart() = noReparse("val t = {<caret>a: Int -> }", ' ')
-    fun testNoArrow() = noReparse("val t = { <caret> }", 'a')
-
-    fun testAfterRemovingParameterComma() = reparse(inIf("{t,<caret>}"), BACKSPACE_FAKE_CHAR)
-    fun testAfterRemovingNoParameterComma() = noReparse(inIf("{,<caret>}"), BACKSPACE_FAKE_CHAR)
-    fun testAfterRemovingNotLastParameterComma() = noReparse(inIf("{a, b,<caret>}"), BACKSPACE_FAKE_CHAR)
-    fun testAfterRemovingSecondParameter() = noReparse(inIf("{a,b<caret>}"), BACKSPACE_FAKE_CHAR)
-    fun testAfterFirstParameterRenamed() = noReparse(inIf("{a<caret>,}"), 'b')
-
-    fun testAfterRemovingFirstParameterWithOther() = reparse(inIf("{a<caret>,b}"), BACKSPACE_FAKE_CHAR)
-    fun testAfterRemovingFirstParameter() = reparse(inIf("{a<caret>,}"), BACKSPACE_FAKE_CHAR)
-    fun testAfterTypeComma() = reparse(inIf("{a<caret>}"), ',')
-
-    fun reparse(text: String, char: Char): Unit = doTest(text, char, true)
-    fun noReparse(text: String, char: Char): Unit = doTest(text, char, false)
-
-    fun doTest(text: String, char: Char, reparse: Boolean) {
-        val file = myFixture.configureByText("a.kt", text.trimMargin())
-
-        val lambdaExpressionBefore = PsiTreeUtil.findChildOfType(file, KtLambdaExpression::class.java)
-
-        performTypingAction(myFixture.editor, char)
-        PsiDocumentManager.getInstance(LightPlatformTestCase.getProject()).commitDocument(myFixture.getDocument(file))
-
-        val lambdaExpressionAfter = PsiTreeUtil.findChildOfType(file, KtLambdaExpression::class.java)
-
-        val actualReparse = lambdaExpressionAfter != lambdaExpressionBefore
-
-        Assert.assertEquals("Lazy element behaviour was unexpected", reparse, actualReparse)
-    }
-
-    private fun inIf(lambda: String) = "val t: Unit = if (true) $lambda else {}"
-
-    override fun getProjectDescriptor() = LightProjectDescriptor.EMPTY_PROJECT_DESCRIPTOR
-}

idea/tests/org/jetbrains/kotlin/idea/editor/LazyElementTypeTestBase.kt

@@ -0,0 +1,41 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.editor
+
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.psi.PsiElement
+import com.intellij.psi.util.PsiTreeUtil
+import com.intellij.testFramework.EditorTestUtil.*
+import com.intellij.testFramework.LightPlatformTestCase
+import com.intellij.testFramework.LightProjectDescriptor
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCaseBase
+import org.junit.Assert
+
+abstract class LazyElementTypeTestBase<T>(private val lazyElementClass: Class<T>) :
+    KotlinLightCodeInsightFixtureTestCaseBase() where T : PsiElement {
+
+    protected fun reparse(text: String, char: Char): Unit = doTest(text, char, true)
+    protected fun noReparse(text: String, char: Char): Unit = doTest(text, char, false)
+
+    fun doTest(text: String, char: Char, reparse: Boolean) {
+        val file = myFixture.configureByText("a.kt", text.trimMargin())
+
+        val expressionBefore = PsiTreeUtil.findChildOfType(file, lazyElementClass)
+
+        performTypingAction(myFixture.editor, char)
+        PsiDocumentManager.getInstance(LightPlatformTestCase.getProject()).commitDocument(myFixture.getDocument(file))
+
+        val expressionAfter = PsiTreeUtil.findChildOfType(file, lazyElementClass)
+
+        val actualReparse = expressionAfter != expressionBefore
+
+        Assert.assertEquals("Lazy element behaviour was unexpected", reparse, actualReparse)
+    }
+
+    protected fun inIf(then: String) = "val t: Unit = if (true) $then else {}"
+
+    override fun getProjectDescriptor() = LightProjectDescriptor.EMPTY_PROJECT_DESCRIPTOR
+}

idea/tests/org/jetbrains/kotlin/idea/editor/LazyKtBlockExpressionTest.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.editor
+
+import com.intellij.testFramework.EditorTestUtil
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.test.JUnit3WithIdeaConfigurationRunner
+import org.junit.runner.RunWith
+
+@RunWith(JUnit3WithIdeaConfigurationRunner::class)
+class LazyKtBlockExpressionTest : LazyElementTypeTestBase<KtBlockExpression>(KtBlockExpression::class.java) {
+
+    fun testSimpleReparse() = noReparse(inIf(" { a<caret>b }"), 'c')
+
+    fun testSimpleNotReparse() = reparse(inIf(" { a-<caret>b }"), '>')
+
+    fun testSplitArrow() = reparse(inIf("{ a: Int -<caret>> }"), ' ')
+
+    fun testDeleteArrow() = reparse(inIf("{ a: Int -<caret>> }"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+
+    fun testImbalance1() = noReparse(inIf(" { {}<caret> }"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+
+    fun testImbalance2() = noReparse(inIf(" { }<caret>"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+
+    fun testBlockWithArrowInside() = noReparse(inIf(" { { a: Int -<caret> a } }"), '>')
+
+    fun testBlockWithArrowInside2() = noReparse(inIf(" { } a: Int -<caret> a } }"), '>')
+
+    fun testBlockWithCommaAsLambdaArgument() = reparse(inIf(" { a<caret>  }"), ',')
+
+    fun testBlockWithCommaAsBlockContent() = noReparse(inIf(" { a.b<caret> c.d = 3  }"), ',')
+
+    fun testBlockWithCommaAsBlockContent2() = noReparse(inIf(" { a(b<caret> c)  }"), ',')
+}

idea/tests/org/jetbrains/kotlin/idea/editor/LazyKtLambdaExpressionTest.kt

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.idea.editor
+
+import com.intellij.testFramework.EditorTestUtil
+import org.jetbrains.kotlin.psi.KtLambdaExpression
+import org.jetbrains.kotlin.test.JUnit3WithIdeaConfigurationRunner
+import org.junit.runner.RunWith
+
+@RunWith(JUnit3WithIdeaConfigurationRunner::class)
+class LazyKtLambdaExpressionTest : LazyElementTypeTestBase<KtLambdaExpression>(KtLambdaExpression::class.java) {
+    fun testSplitArrow() = reparse("val t = { a: Int -<caret>> }", ' ')
+    fun testDeleteArrow() = reparse("val t = { a: Int -><caret> }", EditorTestUtil.BACKSPACE_FAKE_CHAR)
+
+    fun testReformatNearArrow() = noReparse("val t = { a: Int<caret>-> }", ' ')
+    fun testChangeAfterArrow() = noReparse("val t = { a: Int -> <caret> }", 'a')
+    fun testDeleteIrrelevantArrow() = noReparse(
+        "val t = { a: Int -> (1..3).filter { b -><caret> b > 2 } }",
+        EditorTestUtil.BACKSPACE_FAKE_CHAR
+    )
+
+    fun testReformatNearLambdaStart() = noReparse("val t = {<caret>a: Int -> }", ' ')
+    fun testNoArrow() = noReparse("val t = { <caret> }", 'a')
+
+    fun testAfterRemovingParameterComma() = reparse(inIf("{t,<caret>}"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+    fun testAfterRemovingNoParameterComma() = noReparse(inIf("{,<caret>}"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+    fun testAfterRemovingNotLastParameterComma() = noReparse(inIf("{a, b,<caret>}"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+    fun testAfterRemovingSecondParameter() = noReparse(inIf("{a,b<caret>}"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+    fun testAfterFirstParameterRenamed() = noReparse(inIf("{a<caret>,}"), 'b')
+
+    fun testAfterRemovingFirstParameterWithOther() = reparse(inIf("{a<caret>,b}"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+    fun testAfterRemovingFirstParameter() = reparse(inIf("{a<caret>,}"), EditorTestUtil.BACKSPACE_FAKE_CHAR)
+    fun testAfterTypeComma() = reparse(inIf("{a<caret>}"), ',')
+}