Uast: implicit returns for local function bodies (KT-32370)

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isPropertyParameter
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.uast.*
+import org.jetbrains.uast.kotlin.expressions.KotlinLocalFunctionULambdaExpression
 import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
 import org.jetbrains.uast.kotlin.internal.KotlinUElementWithComments
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
@@ -188,7 +189,7 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
         }
     }
 
-    if (result is UMethod
+    if ((result is UMethod || result is KotlinLocalFunctionULambdaExpression)
         && result !is KotlinConstructorUMethod // no sense to wrap super calls with `return`
         && element is UExpression
         && element !is UBlockExpression

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/KotlinAbstractUElement.kt.183

@@ -28,6 +28,7 @@ import org.jetbrains.kotlin.psi.psiUtil.getStrictParentOfType
 import org.jetbrains.kotlin.psi.psiUtil.isPropertyParameter
 import org.jetbrains.kotlin.psi.psiUtil.parentsWithSelf
 import org.jetbrains.uast.*
+import org.jetbrains.uast.kotlin.expressions.KotlinLocalFunctionULambdaExpression
 import org.jetbrains.uast.kotlin.expressions.KotlinUElvisExpression
 import org.jetbrains.uast.kotlin.internal.KotlinUElementWithComments
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
@@ -184,7 +185,7 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
         }
     }
 
-    if (result is UMethod
+    if ((result is UMethod || result is KotlinLocalFunctionULambdaExpression)
         && result !is KotlinConstructorUMethod // no sense to wrap super calls with `return`
         && element is UExpression
         && element !is UBlockExpression

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt

@@ -102,18 +102,7 @@ open class KotlinUMethod(
             else -> null
         } ?: return@lz null
 
-        when (bodyExpression) {
+        wrapExpressionBody(this, bodyExpression)
-            !is KtBlockExpression -> {
-                KotlinUBlockExpression.KotlinLazyUBlockExpression(this, { block ->
-                    val implicitReturn = KotlinUImplicitReturnExpression(block)
-                    val uBody = getLanguagePlugin().convertElement(bodyExpression, implicitReturn) as? UExpression
-                            ?: return@KotlinLazyUBlockExpression emptyList()
-                    listOf(implicitReturn.apply { returnExpression = uBody })
-                })
-
-            }
-            else -> getLanguagePlugin().convertElement(bodyExpression, this) as? UExpression
-        }
     }
 
     override val isOverride: Boolean
@@ -142,4 +131,17 @@ open class KotlinUMethod(
                 else
                     KotlinUMethod(psi, containingElement)
     }
-}
+}
+
+internal fun wrapExpressionBody(function: UElement, bodyExpression: KtExpression): UExpression? = when (bodyExpression) {
+    !is KtBlockExpression -> {
+        KotlinUBlockExpression.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
+}

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt.183

@@ -102,18 +102,7 @@ open class KotlinUMethod(
             else -> null
         } ?: return@lz null
 
-        when (bodyExpression) {
+        wrapExpressionBody(this, bodyExpression)
-            !is KtBlockExpression -> {
-                KotlinUBlockExpression.KotlinLazyUBlockExpression(this, { block ->
-                    val implicitReturn = KotlinUImplicitReturnExpression(block)
-                    val uBody = getLanguagePlugin().convertElement(bodyExpression, implicitReturn) as? UExpression
-                        ?: return@KotlinLazyUBlockExpression emptyList()
-                    listOf(implicitReturn.apply { returnExpression = uBody })
-                })
-
-            }
-            else -> getLanguagePlugin().convertElement(bodyExpression, this) as? UExpression
-        }
     }
 
     override val isOverride: Boolean
@@ -136,4 +125,17 @@ open class KotlinUMethod(
                 else
                     KotlinUMethod(psi, containingElement)
     }
-}
+}
+
+internal fun wrapExpressionBody(function: UElement, bodyExpression: KtExpression): UExpression? = when (bodyExpression) {
+    !is KtBlockExpression -> {
+        KotlinUBlockExpression.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
+}

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/declarations/KotlinUMethod.kt.192

@@ -102,18 +102,7 @@ open class KotlinUMethod(
             else -> null
         } ?: return@lz null
 
-        when (bodyExpression) {
+        wrapExpressionBody(this, bodyExpression)
-            !is KtBlockExpression -> {
-                KotlinUBlockExpression.KotlinLazyUBlockExpression(this, { block ->
-                    val implicitReturn = KotlinUImplicitReturnExpression(block)
-                    val uBody = getLanguagePlugin().convertElement(bodyExpression, implicitReturn) as? UExpression
-                            ?: return@KotlinLazyUBlockExpression emptyList()
-                    listOf(implicitReturn.apply { returnExpression = uBody })
-                })
-
-            }
-            else -> getLanguagePlugin().convertElement(bodyExpression, this) as? UExpression
-        }
     }
 
     override fun getBody(): PsiCodeBlock? = super<UAnnotationMethod>.getBody()
@@ -139,4 +128,17 @@ open class KotlinUMethod(
                 else
                     KotlinUMethod(psi, containingElement)
     }
-}
+}
+
+internal fun wrapExpressionBody(function: UElement, bodyExpression: KtExpression): UExpression? = when (bodyExpression) {
+    !is KtBlockExpression -> {
+        KotlinUBlockExpression.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
+}

plugins/uast-kotlin/src/org/jetbrains/uast/kotlin/expressions/LocalFunction.kt

@@ -7,6 +7,7 @@ import com.intellij.psi.PsiVariable
 import org.jetbrains.kotlin.psi.KtFunction
 import org.jetbrains.uast.*
 import org.jetbrains.uast.kotlin.*
+import org.jetbrains.uast.kotlin.declarations.wrapExpressionBody
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiParameter
 import org.jetbrains.uast.kotlin.psi.UastKotlinPsiVariable
 
@@ -36,13 +37,15 @@ internal class KotlinLocalFunctionUVariable(
 }
 
 
-private class KotlinLocalFunctionULambdaExpression(
+internal class KotlinLocalFunctionULambdaExpression(
         override val sourcePsi: KtFunction,
         givenParent: UElement?
 ): KotlinAbstractUExpression(givenParent), ULambdaExpression {
     override val functionalInterfaceType: PsiType? = null
 
-    override val body by lz { KotlinConverter.convertOrEmpty(sourcePsi.bodyExpression, this) }
+    override val body by lz {
+        sourcePsi.bodyExpression?.let { wrapExpressionBody(this, it) } ?: UastEmptyExpression(this)
+    }
 
     override val valueParameters by lz {
         sourcePsi.valueParameters.mapIndexed { i, p ->

plugins/uast-kotlin/testData/LambdaReturn.kt

@@ -29,6 +29,7 @@ fun foo() {
             return "2"
     }
 
+    val lam5 = fun(a: Int) = "a" + a
 
     bar {
         if (it > 5) return

plugins/uast-kotlin/testData/LambdaReturn.log.txt

@@ -69,6 +69,16 @@ UFile (package = org.jetbrains.uast.kotlin)
                                         ULiteralExpression (value = "1")
                                     UReturnExpression
                                         ULiteralExpression (value = "2")
+                UDeclarationsExpression
+                    ULocalVariable (name = lam5)
+                        ULambdaExpression
+                            UParameter (name = a)
+                                UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+                            UBlockExpression
+                                UReturnExpression
+                                    UBinaryExpression (operator = +)
+                                        ULiteralExpression (value = "a")
+                                        USimpleNameReferenceExpression (identifier = a)
                 UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))
                     UIdentifier (Identifier (bar))
                     USimpleNameReferenceExpression (identifier = bar, resolvesTo = null)

plugins/uast-kotlin/testData/LambdaReturn.render.txt

@@ -21,6 +21,9 @@ public final class LambdaReturnKt {
                 if (a < 5) return "5"
                 if (a > 0) return "1" else return "2"
         }
+        var lam5: kotlin.jvm.functions.Function1<? super java.lang.Integer,? extends java.lang.String> = fun (@org.jetbrains.annotations.NotNull var a: int) {
+                return "a" + a
+        }
         bar({
             if (it > 5) return
             var b: int = 1

plugins/uast-kotlin/testData/LocalDeclarations.log.txt

@@ -8,9 +8,11 @@ UFile (package = )
                 UDeclarationsExpression
                     UVariable (name = bar)
                         ULambdaExpression
-                            UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))
+                            UBlockExpression
-                                UIdentifier (Identifier (Local))
+                                UReturnExpression
-                                USimpleNameReferenceExpression (identifier = <init>, resolvesTo = Local)
+                                    UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0))
+                                        UIdentifier (Identifier (Local))
+                                        USimpleNameReferenceExpression (identifier = <init>, resolvesTo = Local)
                 UDeclarationsExpression
                     ULocalVariable (name = baz)
                         ULambdaExpression
@@ -23,7 +25,9 @@ UFile (package = )
                         ULambdaExpression
                             UParameter (name = text)
                                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
-                            ULiteralExpression (value = 42)
+                            UBlockExpression
+                                UReturnExpression
+                                    ULiteralExpression (value = 42)
                 UDeclarationsExpression
                     UClass (name = LocalObject)
                         UField (name = INSTANCE)

plugins/uast-kotlin/testData/LocalDeclarations.render.txt

@@ -4,13 +4,13 @@ public final class LocalDeclarationsKt {
             public fun LocalDeclarationsKt$foo$Local() = UastEmptyExpression
         }
         var bar: <ErrorType> = fun () {
-            <init>()
+                return <init>()
         }
         var baz: kotlin.jvm.functions.Function0<? extends kotlin.Unit> = fun () {
                 <init>()
         }
         var someLocalFun: kotlin.jvm.functions.Function2<? super java.lang.Integer,? super java.lang.String,? extends java.lang.Integer> = fun (@org.jetbrains.annotations.NotNull var text: java.lang.String) {
-            42
+                return 42
         }
         public static final class LocalObject {
             @null public static final var INSTANCE: LocalDeclarationsKt.foo.LocalObject

plugins/uast-kotlin/testData/LocalDeclarations.types.txt

@@ -8,9 +8,11 @@ UFile (package = ) [public final class LocalDeclarationsKt {...]
                 UDeclarationsExpression [var bar: <ErrorType> = fun () {...}]
                     UVariable (name = bar) [var bar: <ErrorType> = fun () {...}]
                         ULambdaExpression [fun () {...}]
-                            UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0)) [<init>()] : PsiType:<ErrorType>
+                            UBlockExpression [{...}]
-                                UIdentifier (Identifier (Local)) [UIdentifier (Identifier (Local))]
+                                UReturnExpression [return <init>()]
-                                USimpleNameReferenceExpression (identifier = <init>, resolvesTo = Local) [<init>] : PsiType:<ErrorType>
+                                    UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0)) [<init>()] : PsiType:<ErrorType>
+                                        UIdentifier (Identifier (Local)) [UIdentifier (Identifier (Local))]
+                                        USimpleNameReferenceExpression (identifier = <init>, resolvesTo = Local) [<init>] : PsiType:<ErrorType>
                 UDeclarationsExpression [var baz: kotlin.jvm.functions.Function0<? extends kotlin.Unit> = fun () {...}]
                     ULocalVariable (name = baz) [var baz: kotlin.jvm.functions.Function0<? extends kotlin.Unit> = fun () {...}]
                         ULambdaExpression [fun () {...}]
@@ -23,7 +25,9 @@ UFile (package = ) [public final class LocalDeclarationsKt {...]
                         ULambdaExpression [fun (@org.jetbrains.annotations.NotNull var text: java.lang.String) {...}]
                             UParameter (name = text) [@org.jetbrains.annotations.NotNull var text: java.lang.String]
                                 UAnnotation (fqName = org.jetbrains.annotations.NotNull) [@org.jetbrains.annotations.NotNull]
-                            ULiteralExpression (value = 42) [42] : PsiType:int
+                            UBlockExpression [{...}]
+                                UReturnExpression [return 42]
+                                    ULiteralExpression (value = 42) [42] : PsiType:int
                 UDeclarationsExpression [public static final class LocalObject {...}]
                     UClass (name = LocalObject) [public static final class LocalObject {...}]
                         UField (name = INSTANCE) [@null public static final var INSTANCE: LocalDeclarationsKt.foo.LocalObject]

plugins/uast-kotlin/testData/LocalDeclarations.values.txt

@@ -8,9 +8,11 @@ UFile (package = ) [public final class LocalDeclarationsKt {...]
                 UDeclarationsExpression [var bar: <ErrorType> = fun () {...}] = Undetermined
                     UVariable (name = bar) [var bar: <ErrorType> = fun () {...}]
                         ULambdaExpression [fun () {...}] = Undetermined
-                            UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0)) [<init>()] = external <init>()()
+                            UBlockExpression [{...}] = Nothing
-                                UIdentifier (Identifier (Local)) [UIdentifier (Identifier (Local))]
+                                UReturnExpression [return <init>()] = Undetermined
-                                USimpleNameReferenceExpression (identifier = <init>, resolvesTo = Local) [<init>] = external <init>()()
+                                    UCallExpression (kind = UastCallKind(name='constructor_call'), argCount = 0)) [<init>()] = external <init>()()
+                                        UIdentifier (Identifier (Local)) [UIdentifier (Identifier (Local))]
+                                        USimpleNameReferenceExpression (identifier = <init>, resolvesTo = Local) [<init>] = external <init>()()
                 UDeclarationsExpression [var baz: kotlin.jvm.functions.Function0<? extends kotlin.Unit> = fun () {...}] = Undetermined
                     ULocalVariable (name = baz) [var baz: kotlin.jvm.functions.Function0<? extends kotlin.Unit> = fun () {...}]
                         ULambdaExpression [fun () {...}] = Undetermined
@@ -23,7 +25,9 @@ UFile (package = ) [public final class LocalDeclarationsKt {...]
                         ULambdaExpression [fun (@org.jetbrains.annotations.NotNull var text: java.lang.String) {...}] = Undetermined
                             UParameter (name = text) [@org.jetbrains.annotations.NotNull var text: java.lang.String]
                                 UAnnotation (fqName = org.jetbrains.annotations.NotNull) [@org.jetbrains.annotations.NotNull]
-                            ULiteralExpression (value = 42) [42] = 42
+                            UBlockExpression [{...}] = Nothing
+                                UReturnExpression [return 42] = Undetermined
+                                    ULiteralExpression (value = 42) [42] = 42
                 UDeclarationsExpression [public static final class LocalObject {...}] = Undetermined
                     UClass (name = LocalObject) [public static final class LocalObject {...}]
                         UField (name = INSTANCE) [@null public static final var INSTANCE: LocalDeclarationsKt.foo.LocalObject]