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

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,6 +45,8 @@ UFile (package = )
                         UObjectLiteralExpression
                             UClass (name = null)
                                 UAnnotationMethod (name = read)
+                                    UBlockExpression
+                                        UReturnExpression
                                             ULiteralExpression (value = 0)
                                 UAnnotationMethod (name = run)
                                     UBlockExpression

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

@@ -1,6 +1,8 @@
 UFile (package = )
     UClass (name = Foo)
         UAnnotationMethod (name = bar)
+            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,8 +3,12 @@ UFile (package = )
         UAnnotationMethod (name = foo)
             UParameter (name = bar)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            UBlockExpression
+                UReturnExpression
                     ULiteralExpression (value = null)
         UAnnotationMethod (name = bar)
+            UBlockExpression
+                UReturnExpression
                     ULiteralExpression (value = 42)
         UAnnotationMethod (name = baz)
             UBlockExpression

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,6 +5,8 @@ UFile (package = )
             USimpleNameReferenceExpression (identifier = Style)
             UClass (name = null)
                 UAnnotationMethod (name = getExitAnimation)
+                    UBlockExpression
+                        UReturnExpression
                             ULiteralExpression (value = "bar")
                 UAnnotationMethod (name = SHEET)
         UField (name = value)

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,6 +7,8 @@ UFile (package = )
             UField (name = b)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UAnnotationMethod (name = getAPlusB)
+                UBlockExpression
+                    UReturnExpression
                         UBinaryExpression (operator = +)
                             USimpleNameReferenceExpression (identifier = a)
                             USimpleNameReferenceExpression (identifier = b)
@@ -19,5 +21,7 @@ UFile (package = )
                     UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UClass (name = Baz)
                 UAnnotationMethod (name = doNothing)
+                    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,6 +1,8 @@
 UFile (package = )
     UClass (name = PropertyTest)
         UAnnotationMethod (name = getStringRepresentation)
+            UBlockExpression
+                UReturnExpression
                     UQualifiedReferenceExpression
                         UThisExpression (label = null)
                         UCallExpression (kind = UastCallKind(name='method_call'), argCount = 0))

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,6 +4,8 @@ UFile (package = )
             UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             ULiteralExpression (value = "/sdcard")
         UAnnotationMethod (name = getWithSetter)
+            UBlockExpression
+                UReturnExpression
                     USimpleNameReferenceExpression (identifier = field)
         UAnnotationMethod (name = setWithSetter)
             UParameter (name = p)

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,6 +4,8 @@ UFile (package = )
             UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             ULiteralExpression (value = "/sdcard")
         UAnnotationMethod (name = getWithoutSetter)
+            UBlockExpression
+                UReturnExpression
                     USimpleNameReferenceExpression (identifier = field)
         UAnnotationMethod (name = setWithoutSetter)
             UParameter (name = p)

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,8 +10,12 @@ UFile (package = )
         UAnnotationMethod (name = getProp1)
         UAnnotationMethod (name = getProp2)
             UAnnotation (fqName = TestAnnotation)
+            UBlockExpression
+                UReturnExpression
                     ULiteralExpression (value = 0)
         UAnnotationMethod (name = getProp3)
+            UBlockExpression
+                UReturnExpression
                     ULiteralExpression (value = 0)
         UAnnotationMethod (name = setProp3)
             UAnnotation (fqName = TestAnnotation)

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,6 +4,8 @@ UFile (package = )
             UParameter (name = $receiver)
                 UAnnotation (fqName = MyReceiverAnnotation)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            UBlockExpression
+                UReturnExpression
                     UQualifiedReferenceExpression
                         UThisExpression (label = null)
                         USimpleNameReferenceExpression (identifier = length)

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,6 +36,8 @@ UFile (package = )
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
             UAnnotationMethod (name = getB)
             UAnnotationMethod (name = getAPlusB)
+                UBlockExpression
+                    UReturnExpression
                         UBinaryExpression (operator = +)
                             USimpleNameReferenceExpression (identifier = a)
                             USimpleNameReferenceExpression (identifier = b)

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,6 +47,8 @@ UFile (package = )
             UParameter (name = i)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
                 ULiteralExpression (value = 0)
+            UBlockExpression
+                UReturnExpression
                     USimpleNameReferenceExpression (identifier = i)
         UAnnotationMethod (name = foo)
             UBlockExpression

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,10 +43,12 @@ 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
+            UBlockExpression [{...}]
+                UReturnExpression [return i]
                     USimpleNameReferenceExpression (identifier = i) [i] : PsiType:int
         UAnnotationMethod (name = foo) [public static final fun foo() : void {...}]
             UBlockExpression [{...}] : PsiType:Unit

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

@@ -3,6 +3,8 @@ UFile (package = )
         UAnnotationMethod (name = foo)
             UParameter (name = bar)
                 UAnnotation (fqName = org.jetbrains.annotations.NotNull)
+            UBlockExpression
+                UReturnExpression
                     USwitchExpression
                         USimpleNameReferenceExpression (identifier = bar)
                         UExpressionList (when)

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

@@ -1,5 +1,6 @@
 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 {
+        return switch (bar)  {
             it is java.lang.String -> {
                 bar
                 break
@@ -13,3 +14,4 @@ public final class WhenIsKt {
         }
 
     }
+}