FIR/UAST: commonize block and return expressions

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/expressions/KotlinAbstractUBlockExpression.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2010-2021 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.uast.kotlin
+
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.uast.UBlockExpression
+import org.jetbrains.uast.UElement
+
+abstract class KotlinAbstractUBlockExpression(
+    override val sourcePsi: KtBlockExpression,
+    givenParent: UElement?
+) : KotlinAbstractUExpression(givenParent), UBlockExpression, KotlinUElementWithType {
+    override val expressions by lz {
+        sourcePsi.statements.map { baseResolveProviderService.baseKotlinConverter.convertOrEmpty(it, this) }
+    }
+}

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/expressions/KotlinLazyUBlockExpression.kt

@@ -0,0 +1,32 @@
+/*
+ * Copyright 2010-2021 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.uast.kotlin
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.psi.KtAnonymousInitializer
+import org.jetbrains.uast.*
+
+class KotlinLazyUBlockExpression(
+    override val uastParent: UElement?,
+    expressionProducer: (expressionParent: UElement) -> List<UExpression>
+) : UBlockExpression {
+    override val psi: PsiElement? get() = null
+    override val javaPsi: PsiElement? get() = null
+    override val sourcePsi: PsiElement? get() = null
+    override val uAnnotations: List<UAnnotation> = emptyList()
+    override val expressions by lz { expressionProducer(this) }
+
+    companion object {
+        fun create(initializers: List<KtAnonymousInitializer>, uastParent: UElement): UBlockExpression {
+            val languagePlugin = uastParent.getLanguagePlugin()
+            return KotlinLazyUBlockExpression(uastParent) { expressionParent ->
+                initializers.map {
+                    languagePlugin.convertOpt(it.body, expressionParent) ?: UastEmptyExpression(expressionParent)
+                }
+            }
+        }
+    }
+}

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/expressions/KotlinUImplicitReturnExpression.kt

@@ -0,0 +1,30 @@
+/*
+ * Copyright 2010-2021 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.uast.kotlin
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UExpression
+import org.jetbrains.uast.UReturnExpression
+import org.jetbrains.uast.kotlin.internal.KotlinFakeUElement
+
+class KotlinUImplicitReturnExpression(
+    givenParent: UElement?
+) : KotlinAbstractUExpression(givenParent), UReturnExpression, KotlinUElementWithType, KotlinFakeUElement {
+    override val psi: PsiElement?
+        get() = null
+
+    override lateinit var returnExpression: UExpression
+
+    // Due to the lack of [psi], (lazily) delegate to the one in [returnExpression]
+    override val baseResolveProviderService: BaseKotlinUastResolveProviderService by lz {
+        (returnExpression as KotlinAbstractUElement).baseResolveProviderService
+    }
+
+    override fun unwrapToSourcePsi(): List<PsiElement> {
+        return returnExpression.toSourcePsiFakeAware()
+    }
+}

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/expressions/KotlinUReturnExpression.kt

@@ -0,0 +1,19 @@
+/*
+ * Copyright 2010-2021 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.uast.kotlin
+
+import org.jetbrains.kotlin.psi.KtReturnExpression
+import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UReturnExpression
+
+class KotlinUReturnExpression(
+    override val sourcePsi: KtReturnExpression,
+    givenParent: UElement?
+) : KotlinAbstractUExpression(givenParent), UReturnExpression, KotlinUElementWithType {
+    override val returnExpression by lz {
+        baseResolveProviderService.baseKotlinConverter.convertOrNull(sourcePsi.returnedExpression, this)
+    }
+}

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/internal/baseKotlinInternalUastUtils.kt

@@ -10,10 +10,7 @@ import com.intellij.psi.PsiFile
 import org.jetbrains.kotlin.asJava.elements.FakeFileForLightClass
 import org.jetbrains.kotlin.asJava.elements.KtLightElement
 import org.jetbrains.kotlin.asJava.elements.KtLightMember
-import org.jetbrains.kotlin.psi.KtBlockExpression
-import org.jetbrains.kotlin.psi.KtDeclaration
-import org.jetbrains.kotlin.psi.KtExpression
-import org.jetbrains.kotlin.psi.KtPsiUtil
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.uast.UDeclaration
 import org.jetbrains.uast.UElement
 import java.util.function.Supplier
@@ -45,3 +42,10 @@ fun KtExpression.unwrapBlockOrParenthesis(): KtExpression {
     }
     return innerExpression
 }
+
+fun KtElement.canAnalyze(): Boolean {
+    if (!isValid) return false
+    val containingFile = containingFile as? KtFile ?: return false // EA-114080, EA-113475, EA-134193
+    if (containingFile.doNotAnalyze != null) return false // To prevent exceptions during analysis
+    return true
+}

plugins/uast-kotlin-base/src/org/jetbrains/uast/kotlin/internal/converterUtils.kt

@@ -7,10 +7,9 @@ package org.jetbrains.uast.kotlin
 
 import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.lexer.KtTokens
-import org.jetbrains.uast.DEFAULT_EXPRESSION_TYPES_LIST
-import org.jetbrains.uast.DEFAULT_TYPES_LIST
-import org.jetbrains.uast.UElement
-import org.jetbrains.uast.UExpression
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.internal.KotlinFakeUElement
 
 fun expressionTypes(requiredType: Class<out UElement>?) =
@@ -54,3 +53,16 @@ fun UElement.toSourcePsiFakeAware(): List<PsiElement> {
     if (this is KotlinFakeUElement) return this.unwrapToSourcePsi()
     return listOfNotNull(this.sourcePsi)
 }
+
+fun wrapExpressionBody(function: UElement, bodyExpression: KtExpression): UExpression? = when (bodyExpression) {
+    !is KtBlockExpression -> {
+        KotlinLazyUBlockExpression(function) { block ->
+            val implicitReturn = KotlinUImplicitReturnExpression(block)
+            val uBody = function.getLanguagePlugin().convertElement(bodyExpression, implicitReturn) as? UExpression
+                ?: return@KotlinLazyUBlockExpression emptyList()
+            listOf(implicitReturn.apply { returnExpression = uBody })
+        }
+
+    }
+    else -> function.getLanguagePlugin().convertElement(bodyExpression, function) as? UExpression
+}