FIR IDE: introduce return statement target provider

generators/tests/org/jetbrains/kotlin/generators/tests/GenerateTests.kt

@@ -88,6 +88,7 @@ import org.jetbrains.kotlin.idea.fir.low.level.api.file.structure.AbstractFileSt
 import org.jetbrains.kotlin.idea.fir.low.level.api.sessions.AbstractSessionsInvalidationTest
 import org.jetbrains.kotlin.idea.fir.low.level.api.trackers.AbstractProjectWideOutOfBlockKotlinModificationTrackerTest
 import org.jetbrains.kotlin.idea.folding.AbstractKotlinFoldingTest
+import org.jetbrains.kotlin.idea.frontend.api.components.AbstractReturnExpressionTargetTest
 import org.jetbrains.kotlin.idea.frontend.api.fir.AbstractResolveCallTest
 import org.jetbrains.kotlin.idea.frontend.api.scopes.AbstractMemberScopeByFqNameTest
 import org.jetbrains.kotlin.idea.frontend.api.symbols.*
@@ -1021,6 +1022,10 @@ fun main(args: Array<String>) {
             testClass<AbstractMemoryLeakInSymbolsTest> {
                 model("symbolMemoryLeak")
             }
+
+            testClass<AbstractReturnExpressionTargetTest> {
+                model("components/returnExpressionTarget")
+            }
         }
 
         testGroup("idea/idea-frontend-fir/idea-fir-low-level-api/tests", "idea/testData") {

idea/idea-fir/src/org/jetbrains/kotlin/idea/completion/KotlinFirCompletionContributor.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.KtSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.KtNamedSymbol
 import org.jetbrains.kotlin.idea.frontend.api.symbols.markers.isExtension
+import org.jetbrains.kotlin.idea.frontend.api.types.KtType
 import org.jetbrains.kotlin.idea.references.mainReference
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.psi.*
@@ -148,3 +149,13 @@ private class KotlinAvailableScopesCompletionProvider(prefixMatcher: PrefixMatch
         collectTypesCompletion(result, implicitScopes)
     }
 }
+
+private object ExpectedTypeProvider {
+    fun KtAnalysisSession.getExpectedType(nameExpression: KtSimpleNameExpression, parameters: CompletionParameters): KtType? =
+        getExpectedTypeByReturnExpression(nameExpression)
+
+    private fun KtAnalysisSession.getExpectedTypeByReturnExpression(nameExpression: KtSimpleNameExpression): KtType? {
+        val parentReturn = nameExpression.parent as? KtReturnExpression ?: return null
+        TODO()
+    }
+}

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/KtAnalysisSession.kt

@@ -45,6 +45,7 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
     protected abstract val symbolDeclarationOverridesProvider: KtSymbolDeclarationOverridesProvider
     @Suppress("LeakingThis")
     protected open val typeRenderer: KtTypeRenderer = KtDefaultTypeRenderer(this, token)
+    protected abstract val expressionHandlingComponent: KtExpressionHandlingComponent
 
     /// TODO: get rid of
     @Deprecated("Used only in completion now, temporary")
@@ -148,4 +149,7 @@ abstract class KtAnalysisSession(final override val token: ValidityToken) : Vali
 
     fun KtType.render(options: KtTypeRendererOptions = KtTypeRendererOptions.DEFAULT): String =
         typeRenderer.render(this, options)
+
+    fun KtReturnExpression.getReturnTargetSymbol(): KtFunctionLikeSymbol? =
+        expressionHandlingComponent.getReturnExpressionTargetSymbol(this)
 }

idea/idea-frontend-api/src/org/jetbrains/kotlin/idea/frontend/api/components/KtExpressionHandlingComponent.kt

@@ -0,0 +1,13 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.frontend.api.components
+
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionLikeSymbol
+import org.jetbrains.kotlin.psi.KtReturnExpression
+
+abstract class KtExpressionHandlingComponent : KtAnalysisSessionComponent() {
+   abstract fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtFunctionLikeSymbol?
+}

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/KtFirAnalysisSession.kt

@@ -45,6 +45,8 @@ private constructor(
     override val symbolDeclarationOverridesProvider: KtSymbolDeclarationOverridesProvider =
         KtFirSymbolDeclarationOverridesProvider(this, token)
 
+    override val expressionHandlingComponent: KtExpressionHandlingComponent = KtFirExpressionHandlingComponent(this, token)
+
     override fun createContextDependentCopy(): KtAnalysisSession {
         check(!isContextSession) { "Cannot create context-dependent copy of KtAnalysis session from a context dependent one" }
         val contextResolveState = LowLevelFirApiFacadeForCompletion.getResolveStateForCompletion(firResolveState)

idea/idea-frontend-fir/src/org/jetbrains/kotlin/idea/frontend/api/fir/components/KtFirExpressionHandlingComponent.kt

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.frontend.api.fir.components
+
+import org.jetbrains.kotlin.fir.expressions.FirReturnExpression
+import org.jetbrains.kotlin.idea.fir.low.level.api.api.getOrBuildFirSafe
+import org.jetbrains.kotlin.idea.frontend.api.ValidityToken
+import org.jetbrains.kotlin.idea.frontend.api.components.KtExpressionHandlingComponent
+import org.jetbrains.kotlin.idea.frontend.api.fir.KtFirAnalysisSession
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtCallableSymbol
+import org.jetbrains.kotlin.idea.frontend.api.symbols.KtFunctionLikeSymbol
+import org.jetbrains.kotlin.psi.KtReturnExpression
+
+internal class KtFirExpressionHandlingComponent(
+    override val analysisSession: KtFirAnalysisSession,
+    override val token: ValidityToken,
+) : KtExpressionHandlingComponent(), KtFirAnalysisSessionComponent {
+    override fun getReturnExpressionTargetSymbol(returnExpression: KtReturnExpression): KtFunctionLikeSymbol? {
+        val fir = returnExpression.getOrBuildFirSafe<FirReturnExpression>(firResolveState) ?: return null
+        val firTargetSymbol = fir.target.labeledElement
+        return firSymbolBuilder.buildCallableSymbol(firTargetSymbol) as KtFunctionLikeSymbol
+    }
+}

idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromFunctionViaLambdaWithoutLabel.kt

@@ -0,0 +1,8 @@
+fun /* EXPECTED_TARGET */x(): Int {
+    receiveLambda {
+        return<caret> 1
+    }
+    return 2
+}
+
+inline fun receiveLambda(x: () -> Unit){}

idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromFunctionWithLabel.kt

@@ -0,0 +1,3 @@
+fun /* EXPECTED_TARGET */x(): Int {
+    return<caret> 1
+}

idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromFunctionWithoutLabel.kt

@@ -0,0 +1,3 @@
+fun /* EXPECTED_TARGET */x(): Int {
+    return<caret> 1
+}

idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromLambdaWithLabel.kt

@@ -0,0 +1,8 @@
+fun x(): Int {
+    receiveLambda { /* EXPECTED_TARGET */
+        return@receiveLambda<caret> 1
+    }
+    return 2
+}
+
+fun receiveLambda(x: () -> Int){}

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/components/AbstractReturnExpressionTargetTest.kt

@@ -0,0 +1,72 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.frontend.api.components
+
+import com.intellij.openapi.util.io.FileUtil
+import com.intellij.psi.PsiComment
+import com.intellij.psi.util.parentOfType
+import junit.framework.Assert
+import org.jetbrains.kotlin.idea.executeOnPooledThreadInReadAction
+import org.jetbrains.kotlin.idea.frontend.api.analyze
+import org.jetbrains.kotlin.idea.test.KotlinLightCodeInsightFixtureTestCase
+import org.jetbrains.kotlin.idea.util.application.executeOnPooledThread
+import org.jetbrains.kotlin.idea.util.application.runReadAction
+import org.jetbrains.kotlin.psi.KtDeclaration
+import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtReturnExpression
+import org.jetbrains.kotlin.psi.psiUtil.forEachDescendantOfType
+import org.jetbrains.kotlin.test.InTextDirectivesUtils
+import java.io.File
+
+abstract class AbstractReturnExpressionTargetTest : KotlinLightCodeInsightFixtureTestCase() {
+    override fun isFirPlugin() = true
+
+    protected fun doTest(path: String) {
+        val testDataFile = File(path)
+        val ktFile = myFixture.configureByText(testDataFile.name, FileUtil.loadFile(testDataFile)) as KtFile
+
+        val ktReturnExpressionAtCaret = ktFile.findElementAt(myFixture.caretOffset)?.parentOfType<KtReturnExpression>()
+            ?: error("No element was found at caret or no <caret> is present in the test file")
+
+        val expectedReturnTarget = ktFile.getExpectedReturnTarget()
+
+        val actualReturnTargetPsi: KtDeclaration? = executeOnPooledThreadInReadAction {
+            analyze(ktFile) {
+                val actualReturnTargetSymbol = ktReturnExpressionAtCaret.getReturnTargetSymbol() ?: return@analyze null
+                actualReturnTargetSymbol.psi as KtDeclaration
+            }
+        }
+
+        Assert.assertEquals(expectedReturnTarget?.text, actualReturnTargetPsi?.text)
+    }
+
+    private fun KtFile.getExpectedReturnTarget(): KtDeclaration? {
+        var declaration: KtDeclaration? = null
+        forEachDescendantOfType<PsiComment> { comment ->
+            if (comment.text == EXPECTED_RETURN_TARGET_COMMENT) {
+                if (declaration != null) {
+                    error("More than one $EXPECTED_RETURN_TARGET_COMMENT found")
+                }
+                declaration = comment.parentOfType()
+            }
+        }
+        val noDeclarationExpected = InTextDirectivesUtils.findStringWithPrefixes(text, NO_TARGET_EXPECTED_PREFIX) != null
+        return when {
+            noDeclarationExpected && declaration != null -> {
+                error("$noDeclarationExpected was present together with $EXPECTED_RETURN_TARGET_COMMENT")
+            }
+            !noDeclarationExpected && declaration == null -> {
+                error("No $EXPECTED_RETURN_TARGET_COMMENT present, but $NO_TARGET_EXPECTED_PREFIX is not provided")
+            }
+            else -> declaration
+        }
+    }
+
+    companion object {
+        private const val EXPECTED_RETURN_TARGET_COMMENT = "/* EXPECTED_TARGET */"
+        private const val NO_TARGET_EXPECTED_PREFIX = "// NO_TARGET_EXPECTED"
+    }
+}

idea/idea-frontend-fir/tests/org/jetbrains/kotlin/idea/frontend/api/components/ReturnExpressionTargetTestGenerated.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2010-2020 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * 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.frontend.api.components;
+
+import com.intellij.testFramework.TestDataPath;
+import org.jetbrains.kotlin.test.JUnit3RunnerWithInners;
+import org.jetbrains.kotlin.test.KotlinTestUtils;
+import org.jetbrains.kotlin.test.TestMetadata;
+import org.junit.runner.RunWith;
+
+import java.io.File;
+import java.util.regex.Pattern;
+
+/** This class is generated by {@link org.jetbrains.kotlin.generators.tests.TestsPackage}. DO NOT MODIFY MANUALLY */
+@SuppressWarnings("all")
+@TestMetadata("idea/idea-frontend-fir/testData/components/returnExpressionTarget")
+@TestDataPath("$PROJECT_ROOT")
+@RunWith(JUnit3RunnerWithInners.class)
+public class ReturnExpressionTargetTestGenerated extends AbstractReturnExpressionTargetTest {
+    private void runTest(String testDataFilePath) throws Exception {
+        KotlinTestUtils.runTest(this::doTest, this, testDataFilePath);
+    }
+
+    public void testAllFilesPresentInReturnExpressionTarget() throws Exception {
+        KotlinTestUtils.assertAllTestsPresentByMetadataWithExcluded(this.getClass(), new File("idea/idea-frontend-fir/testData/components/returnExpressionTarget"), Pattern.compile("^(.+)\\.kt$"), null, true);
+    }
+
+    @TestMetadata("returnFromFunctionViaLambdaWithoutLabel.kt")
+    public void testReturnFromFunctionViaLambdaWithoutLabel() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromFunctionViaLambdaWithoutLabel.kt");
+    }
+
+    @TestMetadata("returnFromFunctionWithLabel.kt")
+    public void testReturnFromFunctionWithLabel() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromFunctionWithLabel.kt");
+    }
+
+    @TestMetadata("returnFromFunctionWithoutLabel.kt")
+    public void testReturnFromFunctionWithoutLabel() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromFunctionWithoutLabel.kt");
+    }
+
+    @TestMetadata("returnFromLambdaWithLabel.kt")
+    public void testReturnFromLambdaWithLabel() throws Exception {
+        runTest("idea/idea-frontend-fir/testData/components/returnExpressionTarget/returnFromLambdaWithLabel.kt");
+    }
+}