Uast: wrapping expression bodies with return and code block (#KT-23557)

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

@@ -167,6 +167,17 @@ fun doConvertParent(element: UElement, parent: PsiElement?): UElement? {
         }
     }
 
+    if (result is UMethod
+        && result !is KotlinConstructorUMethod // no sense to wrap super calls with `return`
+        && element is UExpression
+        && element !is UBlockExpression
+        && element !is UTypeReferenceExpression // when element is a type in extension methods
+    ) {
+        return KotlinUBlockExpression.KotlinLazyUBlockExpression(result, { block ->
+            listOf(KotlinUImplicitReturnExpression(block).apply { returnExpression = element })
+        }).expressions.single()
+    }
+
     return result
 }
 

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

@@ -90,7 +90,18 @@ open class KotlinUMethod(
             else -> null
         } ?: return@lz null
 
-        getLanguagePlugin().convertElement(bodyExpression, this) as? UExpression
+        when (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

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

@@ -16,8 +16,10 @@
 
 package org.jetbrains.uast.kotlin
 
+import com.intellij.psi.PsiElement
 import org.jetbrains.kotlin.psi.KtReturnExpression
 import org.jetbrains.uast.UElement
+import org.jetbrains.uast.UExpression
 import org.jetbrains.uast.UReturnExpression
 
 class KotlinUReturnExpression(
@@ -26,3 +28,14 @@ class KotlinUReturnExpression(
 ) : KotlinAbstractUExpression(givenParent), UReturnExpression, KotlinUElementWithType {
     override val returnExpression by lz { KotlinConverter.convertOrNull(psi.returnedExpression, this) }
 }
+
+class KotlinUImplicitReturnExpression(
+    givenParent: UElement?
+) : KotlinAbstractUExpression(givenParent), UReturnExpression, KotlinUElementWithType {
+    override val psi: PsiElement?
+        get() = null
+
+    override lateinit var returnExpression: UExpression
+        internal set
+
+}

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

@@ -18,7 +18,9 @@ UFile (package = )
             UAnnotation (fqName = SuppressLint)
                 UNamedExpression (name = value)
                     ULiteralExpression (value = "Lorem")
-            ULiteralExpression (value = 5)
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = 5)
         UAnnotationMethod (name = bar)
             UAnnotation (fqName = IntRange)
                 UNamedExpression (name = from)
@@ -31,7 +33,9 @@ UFile (package = )
                         ULiteralExpression (value = "Lorem")
                         ULiteralExpression (value = "Ipsum")
                         ULiteralExpression (value = "Dolor")
-            USimpleNameReferenceExpression (identifier = Unit)
+            UBlockExpression
+                UReturnExpression
+                    USimpleNameReferenceExpression (identifier = Unit)
         UAnnotationMethod (name = fooWithArrLiteral)
             UAnnotation (fqName = RequiresPermission)
                 UNamedExpression (name = anyOf)
@@ -40,7 +44,9 @@ UFile (package = )
                         ULiteralExpression (value = 1)
                         ULiteralExpression (value = 2)
                         ULiteralExpression (value = 3)
-            ULiteralExpression (value = 5)
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = 5)
         UAnnotationMethod (name = fooWithStrArrLiteral)
             UAnnotation (fqName = RequiresStrPermission)
                 UNamedExpression (name = strs)
@@ -49,7 +55,9 @@ UFile (package = )
                         ULiteralExpression (value = "a")
                         ULiteralExpression (value = "b")
                         ULiteralExpression (value = "c")
-            ULiteralExpression (value = 3)
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = 3)
     UClass (name = IntRange)
         UAnnotationMethod (name = from)
         UAnnotationMethod (name = to)

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

@@ -3,14 +3,22 @@ public final class AnnotationParametersKt {
     @IntRange(from = 10, to = 0)
     @WithDefaultValue
     @SuppressLint(value = "Lorem")
-    public static final fun foo() : int = 5
+    public static final fun foo() : int {
+        return 5
+    }
     @IntRange(from = 0, to = 100)
     @SuppressLint(value = <noref>("Lorem", "Ipsum", "Dolor"))
-    public static final fun bar() : void = Unit
+    public static final fun bar() : void {
+        return Unit
+    }
     @RequiresPermission(anyOf = collectionLiteral[1, 2, 3])
-    public static final fun fooWithArrLiteral() : int = 5
+    public static final fun fooWithArrLiteral() : int {
+        return 5
+    }
     @RequiresStrPermission(strs = collectionLiteral["a", "b", "c"])
-    public static final fun fooWithStrArrLiteral() : int = 3
+    public static final fun fooWithStrArrLiteral() : int {
+        return 3
+    }
 }
 
 public abstract annotation IntRange {

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

@@ -45,7 +45,9 @@ UFile (package = )
                         UObjectLiteralExpression
                             UClass (name = null)
                                 UAnnotationMethod (name = read)
-                                    ULiteralExpression (value = 0)
+                                    UBlockExpression
+                                        UReturnExpression
+                                            ULiteralExpression (value = 0)
                                 UAnnotationMethod (name = run)
                                     UBlockExpression
                                 UAnnotationMethod (name = AnonymousKt$foo$runnableIs$1)

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

@@ -1,6 +1,8 @@
 UFile (package = )
     UClass (name = Foo)
         UAnnotationMethod (name = bar)
-            ULiteralExpression (value = "Hello!")
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = "Hello!")
     UClass (name = Baz)
         UAnnotationMethod (name = Baz)

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

@@ -1,5 +1,7 @@
 public abstract interface Foo {
-    public default fun bar() : java.lang.String = "Hello!"
+    public default fun bar() : java.lang.String {
+        return "Hello!"
+    }
 }
 
 public final class Baz : Foo {

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

@@ -3,9 +3,13 @@ UFile (package = )
         UAnnotationMethod (name = foo)
             UParameter (name = bar)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
-            ULiteralExpression (value = null)
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = null)
         UAnnotationMethod (name = bar)
-            ULiteralExpression (value = 42)
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = 42)
         UAnnotationMethod (name = baz)
             UBlockExpression
                 UReturnExpression

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

@@ -1,6 +1,10 @@
 public final class ElvisKt {
-    public static final fun foo(@org.jetbrains.annotations.NotNull bar: java.lang.String) : java.lang.String = null
+    public static final fun foo(@org.jetbrains.annotations.NotNull bar: java.lang.String) : java.lang.String {
-    public static final fun bar() : int = 42
+        return null
+    }
+    public static final fun bar() : int {
+        return 42
+    }
     public static final fun baz() : java.lang.String {
         return elvis {
             @null var var243c51a0: java.lang.String = elvis {

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

@@ -5,7 +5,9 @@ UFile (package = )
             USimpleNameReferenceExpression (identifier = Style)
             UClass (name = null)
                 UAnnotationMethod (name = getExitAnimation)
-                    ULiteralExpression (value = "bar")
+                    UBlockExpression
+                        UReturnExpression
+                            ULiteralExpression (value = "bar")
                 UAnnotationMethod (name = SHEET)
         UField (name = value)
             UAnnotation (fqName = org.jetbrains.annotations.NotNull)

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

@@ -1,6 +1,8 @@
 public enum Style {
     @null SHEET {
-        public fun getExitAnimation() : java.lang.String = "bar"
+        public fun getExitAnimation() : java.lang.String {
+            return "bar"
+        }
         fun SHEET() = UastEmptyExpression
     }
     @org.jetbrains.annotations.NotNull private final var value: java.lang.String

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

@@ -7,9 +7,11 @@ UFile (package = )
             UField (name = b)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UAnnotationMethod (name = getAPlusB)
-                UBinaryExpression (operator = +)
+                UBlockExpression
-                    USimpleNameReferenceExpression (identifier = a)
+                    UReturnExpression
-                    USimpleNameReferenceExpression (identifier = b)
+                        UBinaryExpression (operator = +)
+                            USimpleNameReferenceExpression (identifier = a)
+                            USimpleNameReferenceExpression (identifier = b)
             UAnnotationMethod (name = getA)
             UAnnotationMethod (name = getB)
             UAnnotationMethod (name = Bar)
@@ -19,5 +21,7 @@ UFile (package = )
                     UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UClass (name = Baz)
                 UAnnotationMethod (name = doNothing)
-                    USimpleNameReferenceExpression (identifier = Unit)
+                    UBlockExpression
+                        UReturnExpression
+                            USimpleNameReferenceExpression (identifier = Unit)
                 UAnnotationMethod (name = Baz)

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

@@ -3,12 +3,16 @@ public final class Foo {
     public static final class Bar {
         @org.jetbrains.annotations.NotNull private final var a: int
         @org.jetbrains.annotations.NotNull private final var b: int
-        public final fun getAPlusB() : int = a + b
+        public final fun getAPlusB() : int {
+            return a + b
+        }
         public final fun getA() : int = UastEmptyExpression
         public final fun getB() : int = UastEmptyExpression
         public fun Bar(@org.jetbrains.annotations.NotNull a: int, @org.jetbrains.annotations.NotNull b: int) = UastEmptyExpression
         public static final class Baz {
-            public final fun doNothing() : void = Unit
+            public final fun doNothing() : void {
+                return Unit
+            }
             public fun Baz() = UastEmptyExpression
         }
     }

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

@@ -1,11 +1,13 @@
 UFile (package = )
     UClass (name = PropertyTest)
         UAnnotationMethod (name = getStringRepresentation)
-            UQualifiedReferenceExpression
+            UBlockExpression
-                UThisExpression (label = null)
+                UReturnExpression
-                UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                    UQualifiedReferenceExpression
-                    UIdentifier (Identifier (toString))
+                        UThisExpression (label = null)
-                    USimpleNameReferenceExpression (identifier = toString)
+                        UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))
+                            UIdentifier (Identifier (toString))
+                            USimpleNameReferenceExpression (identifier = toString)
         UAnnotationMethod (name = setStringRepresentation)
             UParameter (name = value)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)

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

@@ -1,5 +1,7 @@
 public final class PropertyTest {
-    public final fun getStringRepresentation() : java.lang.String = this.toString()
+    public final fun getStringRepresentation() : java.lang.String {
+        return this.toString()
+    }
     public final fun setStringRepresentation(@org.jetbrains.annotations.NotNull value: java.lang.String) : void {
         setDataFromString(value)
     }

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

@@ -4,7 +4,9 @@ UFile (package = )
             UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             ULiteralExpression (value = "/sdcard")
         UAnnotationMethod (name = getWithSetter)
-            USimpleNameReferenceExpression (identifier = field)
+            UBlockExpression
+                UReturnExpression
+                    USimpleNameReferenceExpression (identifier = field)
         UAnnotationMethod (name = setWithSetter)
             UParameter (name = p)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)

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

@@ -1,6 +1,8 @@
 public final class TestPropertyInitializer {
     @org.jetbrains.annotations.NotNull private var withSetter: java.lang.String = "/sdcard"
-    public final fun getWithSetter() : java.lang.String = field
+    public final fun getWithSetter() : java.lang.String {
+        return field
+    }
     public final fun setWithSetter(@org.jetbrains.annotations.NotNull p: java.lang.String) : void {
         field = p
     }

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

@@ -4,7 +4,9 @@ UFile (package = )
             UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             ULiteralExpression (value = "/sdcard")
         UAnnotationMethod (name = getWithoutSetter)
-            USimpleNameReferenceExpression (identifier = field)
+            UBlockExpression
+                UReturnExpression
+                    USimpleNameReferenceExpression (identifier = field)
         UAnnotationMethod (name = setWithoutSetter)
             UParameter (name = p)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)

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

@@ -1,5 +1,7 @@
 public final class PropertyInitializerWithoutSetterKt {
     @org.jetbrains.annotations.NotNull private static var withoutSetter: java.lang.String = "/sdcard"
-    public static final fun getWithoutSetter() : java.lang.String = field
+    public static final fun getWithoutSetter() : java.lang.String {
+        return field
+    }
     public static final fun setWithoutSetter(@org.jetbrains.annotations.NotNull p: java.lang.String) : void = UastEmptyExpression
 }

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

@@ -10,9 +10,13 @@ UFile (package = )
         UAnnotationMethod (name = getProp1)
         UAnnotationMethod (name = getProp2)
             UAnnotation (fqName = TestAnnotation)
-            ULiteralExpression (value = 0)
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = 0)
         UAnnotationMethod (name = getProp3)
-            ULiteralExpression (value = 0)
+            UBlockExpression
+                UReturnExpression
+                    ULiteralExpression (value = 0)
         UAnnotationMethod (name = setProp3)
             UAnnotation (fqName = TestAnnotation)
             UParameter (name = value)

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

@@ -3,8 +3,12 @@ public final class PropertyWithAnnotationKt {
     @org.jetbrains.annotations.NotNull private static var prop3: int = 0
     public static final fun getProp1() : int = UastEmptyExpression
     @TestAnnotation
-    public static final fun getProp2() : int = 0
+    public static final fun getProp2() : int {
-    public static final fun getProp3() : int = 0
+        return 0
+    }
+    public static final fun getProp3() : int {
+        return 0
+    }
     @TestAnnotation
     public static final fun setProp3(@org.jetbrains.annotations.NotNull value: int) : void {
         field = value

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

@@ -4,7 +4,9 @@ UFile (package = )
             UParameter (name = $receiver)
                 UAnnotation (fqName = MyReceiverAnnotation)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
-            UQualifiedReferenceExpression
+            UBlockExpression
-                UThisExpression (label = null)
+                UReturnExpression
-                USimpleNameReferenceExpression (identifier = length)
+                    UQualifiedReferenceExpression
+                        UThisExpression (label = null)
+                        USimpleNameReferenceExpression (identifier = length)
     UClass (name = MyReceiverAnnotation)

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

@@ -1,5 +1,7 @@
 public final class ReceiverFunKt {
-    public static final fun foo(@MyReceiverAnnotation @org.jetbrains.annotations.NotNull $receiver: java.lang.String) : int = this.length
+    public static final fun foo(@MyReceiverAnnotation @org.jetbrains.annotations.NotNull $receiver: java.lang.String) : int {
+        return this.length
+    }
 }
 
 public abstract annotation MyReceiverAnnotation {

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

@@ -36,9 +36,11 @@ UFile (package = )
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UAnnotationMethod (name = getB)
             UAnnotationMethod (name = getAPlusB)
-                UBinaryExpression (operator = +)
+                UBlockExpression
-                    USimpleNameReferenceExpression (identifier = a)
+                    UReturnExpression
-                    USimpleNameReferenceExpression (identifier = b)
+                        UBinaryExpression (operator = +)
+                            USimpleNameReferenceExpression (identifier = a)
+                            USimpleNameReferenceExpression (identifier = b)
             UAnnotationMethod (name = getA)
             UAnnotationMethod (name = Bar)
                 UParameter (name = a)

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

@@ -11,7 +11,9 @@ public class SimpleScript {
         @org.jetbrains.annotations.NotNull private final var b: int = 0
         @org.jetbrains.annotations.NotNull private final var a: int
         public final fun getB() : int = UastEmptyExpression
-        public final fun getAPlusB() : int = a + b
+        public final fun getAPlusB() : int {
+            return a + b
+        }
         public final fun getA() : int = UastEmptyExpression
         public fun Bar(@org.jetbrains.annotations.NotNull a: int) = UastEmptyExpression
         public static final class Baz {

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

@@ -47,7 +47,9 @@ UFile (package = )
             UParameter (name = i)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
                 ULiteralExpression (value = 0)
-            USimpleNameReferenceExpression (identifier = i)
+            UBlockExpression
+                UReturnExpression
+                    USimpleNameReferenceExpression (identifier = i)
         UAnnotationMethod (name = foo)
             UBlockExpression
                 UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1))

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

@@ -9,7 +9,9 @@ public final class StringTemplateComplexKt {
     public static final fun getCase5() : java.lang.String = UastEmptyExpression
     public static final fun getLiteralInLiteral() : java.lang.String = UastEmptyExpression
     public static final fun getLiteralInLiteral2() : java.lang.String = UastEmptyExpression
-    public static final fun simpleForTemplate(@org.jetbrains.annotations.NotNull i: int) : java.lang.String = i
+    public static final fun simpleForTemplate(@org.jetbrains.annotations.NotNull i: int) : java.lang.String {
+        return i
+    }
     public static final fun foo() : void {
         println(baz)
         var template1: java.lang.String = simpleForTemplate()

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

@@ -43,11 +43,13 @@ UFile (package = ) [public final class StringTemplateComplexKt {...]
         UAnnotationMethod (name = getCase5) [public static final fun getCase5() : java.lang.String = UastEmptyExpression]
         UAnnotationMethod (name = getLiteralInLiteral) [public static final fun getLiteralInLiteral() : java.lang.String = UastEmptyExpression]
         UAnnotationMethod (name = getLiteralInLiteral2) [public static final fun getLiteralInLiteral2() : java.lang.String = UastEmptyExpression]
-        UAnnotationMethod (name = simpleForTemplate) [public static final fun simpleForTemplate(@org.jetbrains.annotations.NotNull i: int) : java.lang.String = i]
+        UAnnotationMethod (name = simpleForTemplate) [public static final fun simpleForTemplate(@org.jetbrains.annotations.NotNull i: int) : java.lang.String {...}]
             UParameter (name = i) [@org.jetbrains.annotations.NotNull var i: int = 0]
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull) [@org.jetbrains.annotations.NotNull]
                 ULiteralExpression (value = 0) [0] : PsiType:int
-            USimpleNameReferenceExpression (identifier = i) [i] : PsiType:int
+            UBlockExpression [{...}]
+                UReturnExpression [return i]
+                    USimpleNameReferenceExpression (identifier = i) [i] : PsiType:int
         UAnnotationMethod (name = foo) [public static final fun foo() : void {...}]
             UBlockExpression [{...}] : PsiType:Unit
                 UCallExpression (kind = UastCallKind(name='method_call'), argCount = 1)) [println(baz)] : PsiType:Unit

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

@@ -3,20 +3,22 @@ UFile (package = )
         UAnnotationMethod (name = foo)
             UParameter (name = bar)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
-            USwitchExpression
+            UBlockExpression
-                USimpleNameReferenceExpression (identifier = bar)
+                UReturnExpression
-                UExpressionList (when)
+                    USwitchExpression
-                    USwitchClauseExpressionWithBody
+                        USimpleNameReferenceExpression (identifier = bar)
-                        UBinaryExpressionWithType
+                        UExpressionList (when)
-                            USimpleNameReferenceExpression (identifier = it)
+                            USwitchClauseExpressionWithBody
-                            UTypeReferenceExpression (name = java.lang.String)
+                                UBinaryExpressionWithType
-                        UExpressionList (when_entry)
+                                    USimpleNameReferenceExpression (identifier = it)
-                            USimpleNameReferenceExpression (identifier = bar)
+                                    UTypeReferenceExpression (name = java.lang.String)
-                            UBreakExpression (label = null)
+                                UExpressionList (when_entry)
-                    USwitchClauseExpressionWithBody
+                                    USimpleNameReferenceExpression (identifier = bar)
-                        UBinaryExpressionWithType
+                                    UBreakExpression (label = null)
-                            USimpleNameReferenceExpression (identifier = it)
+                            USwitchClauseExpressionWithBody
-                            UTypeReferenceExpression (name = java.lang.String)
+                                UBinaryExpressionWithType
-                        UExpressionList (when_entry)
+                                    USimpleNameReferenceExpression (identifier = it)
-                            ULiteralExpression (value = "<error>")
+                                    UTypeReferenceExpression (name = java.lang.String)
-                            UBreakExpression (label = null)
+                                UExpressionList (when_entry)
+                                    ULiteralExpression (value = "<error>")
+                                    UBreakExpression (label = null)

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

@@ -1,15 +1,17 @@
 public final class WhenIsKt {
-    public static final fun foo(@org.jetbrains.annotations.NotNull bar: java.lang.Object) : java.lang.String = switch (bar)  {
+    public static final fun foo(@org.jetbrains.annotations.NotNull bar: java.lang.Object) : java.lang.String {
-        it is java.lang.String -> {
+        return switch (bar)  {
-            bar
+            it is java.lang.String -> {
-            break
+                bar
-        }
+                break
+            }
+
+            it !is java.lang.String -> {
+                "<error>"
+                break
+            }
 
-        it !is java.lang.String -> {
-            "<error>"
-            break
         }
 
     }
-
 }