Add bodyBlockExpression() helper function to KtDeclarationWithBody

compiler/frontend/src/org/jetbrains/kotlin/diagnostics/PositioningStrategies.kt

@@ -437,7 +437,7 @@ object PositioningStrategies {
     @JvmField
     val DECLARATION_WITH_BODY: PositioningStrategy<KtDeclarationWithBody> = object : PositioningStrategy<KtDeclarationWithBody>() {
         override fun mark(element: KtDeclarationWithBody): List<TextRange> {
-            val lastBracketRange = (element.bodyExpression as? KtBlockExpression)?.lastBracketRange
+            val lastBracketRange = element.bodyBlockExpression?.lastBracketRange
             return if (lastBracketRange != null)
                 markRange(lastBracketRange)
             else
@@ -445,7 +445,7 @@ object PositioningStrategies {
         }
 
         override fun isValid(element: KtDeclarationWithBody): Boolean {
-            return super.isValid(element) && (element.bodyExpression as? KtBlockExpression)?.lastBracketRange != null
+            return super.isValid(element) && element.bodyBlockExpression?.lastBracketRange != null
         }
     }
 

compiler/psi/src/org/jetbrains/kotlin/psi/KtDeclarationWithBody.java

@@ -41,5 +41,15 @@ public interface KtDeclarationWithBody extends KtDeclaration {
 
     @NotNull
     List<KtParameter> getValueParameters();
+
+    @Nullable
+    default KtBlockExpression getBodyBlockExpression() {
+        KtExpression bodyExpression = getBodyExpression();
+        if (bodyExpression instanceof KtBlockExpression) {
+            return (KtBlockExpression) bodyExpression;
+        }
+
+        return null;
+    }
 }
 

compiler/psi/src/org/jetbrains/kotlin/psi/KtPsiFactory.kt

@@ -276,7 +276,7 @@ class KtPsiFactory @JvmOverloads constructor(private val project: Project, val m
     }
 
     fun createDestructuringDeclaration(text: String): KtDestructuringDeclaration {
-        return (createFunction("fun foo() {$text}").bodyExpression as KtBlockExpression).statements.first() as KtDestructuringDeclaration
+        return createFunction("fun foo() {$text}").bodyBlockExpression!!.statements.first() as KtDestructuringDeclaration
     }
 
     fun createDestructuringParameter(text: String): KtParameter {
@@ -337,7 +337,7 @@ class KtPsiFactory @JvmOverloads constructor(private val project: Project, val m
     }
 
     fun createEmptyBody(): KtBlockExpression {
-        return createFunction("fun foo() {}").bodyExpression as KtBlockExpression
+        return createFunction("fun foo() {}").bodyBlockExpression!!
     }
 
     fun createAnonymousInitializer(): KtAnonymousInitializer {
@@ -821,13 +821,13 @@ class KtPsiFactory @JvmOverloads constructor(private val project: Project, val m
     }
 
     fun createBlock(bodyText: String): KtBlockExpression {
-        return createFunction("fun foo() {\n$bodyText\n}").bodyExpression as KtBlockExpression
+        return createFunction("fun foo() {\n$bodyText\n}").bodyBlockExpression!!
     }
 
     fun createSingleStatementBlock(statement: KtExpression, prevComment: String? = null, nextComment: String? = null): KtBlockExpression {
         val prev = if (prevComment == null) "" else " $prevComment "
         val next = if (nextComment == null) "" else " $nextComment "
-        return createDeclarationByPattern<KtNamedFunction>("fun foo() {\n$prev$0$next\n}", statement).bodyExpression as KtBlockExpression
+        return createDeclarationByPattern<KtNamedFunction>("fun foo() {\n$prev$0$next\n}", statement).bodyBlockExpression!!
     }
 
     fun createComment(text: String): PsiComment {

compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt

@@ -35,7 +35,6 @@ import org.jetbrains.kotlin.name.SpecialNames
 import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.stubs.KotlinClassOrObjectStub
 import org.jetbrains.kotlin.types.expressions.OperatorConventions
-import java.lang.IllegalArgumentException
 import java.util.*
 
 // NOTE: in this file we collect only Kotlin-specific methods working with PSI and not modifying it
@@ -300,7 +299,7 @@ fun KtNamedFunction.isContractPresentPsiCheck(): Boolean {
         !hasModifier(KtTokens.OPERATOR_KEYWORD)
     if (!contractAllowedHere) return false
 
-    val firstExpression = ((this as? KtFunction)?.bodyExpression as? KtBlockExpression)?.statements?.firstOrNull() ?: return false
+    val firstExpression = (this as? KtFunction)?.bodyBlockExpression?.statements?.firstOrNull() ?: return false
 
     return firstExpression.isContractDescriptionCallPsiCheck()
 }

idea/idea-jvm/src/org/jetbrains/kotlin/idea/debugger/breakpoints/breakpointTypeUtils.kt

@@ -143,7 +143,7 @@ fun getLambdasAtLineIfAny(file: KtFile, line: Int): List<KtFunction> {
             .toSet()
 
     return allLiterals.filter {
-        val statement = (it.bodyExpression as? KtBlockExpression)?.statements?.firstOrNull() ?: it
+        val statement = it.bodyBlockExpression?.statements?.firstOrNull() ?: it
         statement.getLineNumber() == line && statement.getLineNumber(false) == line
     }
 }

idea/idea-jvm/src/org/jetbrains/kotlin/idea/debugger/evaluate/KotlinCodeFragmentFactory.kt

@@ -144,7 +144,7 @@ class KotlinCodeFragmentFactory : CodeFragmentFactory() {
 
                 val fakeFile = createFakeFileWithJavaContextElement(fakeFunctionText, contextElement)
                 val fakeFunction = fakeFile.declarations.firstOrNull() as? KtFunction
-                val fakeContext = (fakeFunction?.bodyExpression as? KtBlockExpression)?.statements?.lastOrNull()
+                val fakeContext = fakeFunction?.bodyBlockExpression?.statements?.lastOrNull()
 
                 return@putCopyableUserData wrapContextIfNeeded(project, contextElement, fakeContext) ?: emptyFile
             })

idea/idea-jvm/src/org/jetbrains/kotlin/idea/debugger/evaluate/extractFunctionForDebuggerUtil.kt

@@ -335,7 +335,7 @@ private fun findElementBefore(contextElement: PsiElement): PsiElement? {
             wrapInLambdaCall(contextElement.bodyExpression!!)
         }
         contextElement is KtDeclarationWithBody && contextElement.hasBlockBody() -> {
-            val block = contextElement.bodyExpression as KtBlockExpression
+            val block = contextElement.bodyBlockExpression!!
             val last = block.statements.lastOrNull()
             last as? KtReturnExpression ?: block.rBrace
         }

idea/src/org/jetbrains/kotlin/idea/inspections/ConflictingExtensionPropertyInspection.kt

@@ -101,7 +101,7 @@ class ConflictingExtensionPropertyInspection : AbstractKotlinInspection() {
 
     private fun checkGetterBodyIsGetMethodCall(getter: KtPropertyAccessor, getMethod: FunctionDescriptor): Boolean {
         return if (getter.hasBlockBody()) {
-            val statement = (getter.bodyExpression as? KtBlockExpression)?.statements?.singleOrNull() ?: return false
+            val statement = getter.bodyBlockExpression?.statements?.singleOrNull() ?: return false
             (statement as? KtReturnExpression)?.returnedExpression.isGetMethodCall(getMethod)
         }
         else {
@@ -112,7 +112,7 @@ class ConflictingExtensionPropertyInspection : AbstractKotlinInspection() {
     private fun checkSetterBodyIsSetMethodCall(setter: KtPropertyAccessor, setMethod: FunctionDescriptor): Boolean {
         val valueParameterName = setter.valueParameters.singleOrNull()?.nameAsName ?: return false
         if (setter.hasBlockBody()) {
-            val statement = (setter.bodyExpression as? KtBlockExpression)?.statements?.singleOrNull() ?: return false
+            val statement = setter.bodyBlockExpression?.statements?.singleOrNull() ?: return false
             return statement.isSetMethodCall(setMethod, valueParameterName)
         }
         else {

idea/src/org/jetbrains/kotlin/idea/inspections/SetterBackingFieldAssignmentInspection.kt

@@ -21,7 +21,7 @@ class SetterBackingFieldAssignmentInspection : AbstractKotlinInspection(), Clean
     override fun buildVisitor(holder: ProblemsHolder, isOnTheFly: Boolean, session: LocalInspectionToolSession): PsiElementVisitor {
         return propertyAccessorVisitor(fun(accessor) {
             if (!accessor.isSetter) return
-            val bodyExpression = accessor.bodyExpression as? KtBlockExpression ?: return
+            val bodyExpression = accessor.bodyBlockExpression ?: return
 
             val property = accessor.property
             val propertyContext = property.analyze()
@@ -76,7 +76,7 @@ private class AssignBackingFieldFix : LocalQuickFix {
     override fun applyFix(project: Project, descriptor: ProblemDescriptor) {
         val setter = descriptor.psiElement as? KtPropertyAccessor ?: return
         val parameter = setter.valueParameters.firstOrNull() ?: return
-        val bodyExpression = setter.bodyExpression as? KtBlockExpression ?: return
+        val bodyExpression = setter.bodyBlockExpression ?: return
         setter.hasBlockBody()
         bodyExpression.addBefore(
             KtPsiFactory(setter).createExpression("field = ${parameter.text}"),

idea/src/org/jetbrains/kotlin/idea/inspections/UseExpressionBodyInspection.kt

@@ -100,7 +100,7 @@ class UseExpressionBodyInspection(private val convertEmptyToUnit: Boolean) : Abs
 
     private fun KtDeclarationWithBody.blockExpression() = when (this) {
         is KtFunctionLiteral -> null
-        else -> if (!hasBlockBody()) null else bodyExpression as? KtBlockExpression
+        else -> if (!hasBlockBody()) null else bodyBlockExpression
     }
 
     private fun KtBlockExpression.findValueStatement(): KtExpression? {

idea/src/org/jetbrains/kotlin/idea/intentions/ConvertFunctionToPropertyIntention.kt

@@ -88,7 +88,7 @@ class ConvertFunctionToPropertyIntention : SelfTargetingIntention<KtNamedFunctio
                     getterExpression(originalFunction.bodyExpression!!.text, breakLine = originalFunction.typeReference != null)
                 }
                 else {
-                    (originalFunction.bodyExpression as? KtBlockExpression)?.let { body ->
+                    originalFunction.bodyBlockExpression?.let { body ->
                         transform {
                             append("\nget() ")
                             append(body.text)

idea/src/org/jetbrains/kotlin/idea/intentions/UsePropertyAccessSyntaxIntention.kt

@@ -243,7 +243,7 @@ class UsePropertyAccessSyntaxIntention :
         var callToConvert = callExpression
         if (callParent is KtDeclarationWithBody && call == callParent.bodyExpression) {
             ConvertToBlockBodyIntention.convert(callParent)
-            val firstStatement = (callParent.bodyExpression as? KtBlockExpression)?.statements?.first()
+            val firstStatement = callParent.bodyBlockExpression?.statements?.first()
             callToConvert = (firstStatement as? KtQualifiedExpression)?.selectorExpression as? KtCallExpression
                     ?: firstStatement as? KtCallExpression
                     ?: throw IllegalStateException("Unexpected contents of function after conversion: ${callParent.text}")

idea/src/org/jetbrains/kotlin/idea/quickfix/createFromUsage/callableBuilder/generationUtils.kt

@@ -49,7 +49,7 @@ fun setupEditorSelection(editor: Editor, declaration: KtNamedDeclaration) {
 
     when (declaration) {
         is KtNamedFunction, is KtSecondaryConstructor -> {
-            ((declaration as KtFunction).bodyExpression as? KtBlockExpression)?.let {
+            (declaration as KtFunction).bodyBlockExpression?.let {
                 positionBetween(it.lBrace!!, it.rBrace!!)
             }
         }

idea/src/org/jetbrains/kotlin/idea/quickfix/migration/MigrateExternalExtensionFix.kt

@@ -188,7 +188,7 @@ class MigrateExternalExtensionFix(declaration: KtNamedDeclaration)
                 }
 
                 val setterStubProperty = ktPsiFactory.createProperty("val x: Unit set(value) { Unit }")
-                val block = setterStubProperty.setter!!.bodyExpression as KtBlockExpression
+                val block = setterStubProperty.setter!!.bodyBlockExpression!!
                 block.statements.single().replace(setterBody)
                 declaration.add(setterStubProperty.setter!!)
             }

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/extractionEngine/extractableAnalysisUtil.kt

@@ -401,7 +401,7 @@ internal fun ExtractionData.createTemporaryCodeBlock(): KtBlockExpression {
     if (options.extractAsProperty) {
         return ((createTemporaryDeclaration("val = {\n$0\n}\n") as KtProperty).initializer as KtLambdaExpression).bodyExpression!!
     }
-    return (createTemporaryDeclaration("fun() {\n$0\n}\n") as KtNamedFunction).bodyExpression as KtBlockExpression
+    return (createTemporaryDeclaration("fun() {\n$0\n}\n") as KtNamedFunction).bodyBlockExpression!!
 }
 
 private fun KotlinType.collectReferencedTypes(processTypeArguments: Boolean): List<KotlinType> {

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceParameter/KotlinIntroduceParameterDialog.kt

@@ -244,7 +244,7 @@ class KotlinIntroduceParameterDialog private constructor(
 
         project.executeCommand(commandName) {
             fun createLambdaForArgument(function: KtFunction): KtExpression {
-                val statement = (function.bodyExpression as KtBlockExpression).statements.single()
+                val statement = function.bodyBlockExpression!!.statements.single()
                 val space = if (statement.isMultiLine()) "\n" else " "
                 val parameters = function.valueParameters
                 val parametersText = if (parameters.isNotEmpty()) {

idea/src/org/jetbrains/kotlin/idea/refactoring/introduce/introduceVariable/KotlinIntroduceVariableHandler.kt

@@ -312,8 +312,8 @@ object KotlinIntroduceVariableHandler : RefactoringActionHandler {
 
             val newDeclaration = ConvertToBlockBodyIntention.convert(commonContainer)
 
-            val newCommonContainer = (newDeclaration.bodyExpression as KtBlockExpression?)
+            val newCommonContainer = newDeclaration.bodyBlockExpression
-                    .sure { "New body is not found: " + newDeclaration }
+                .sure { "New body is not found: " + newDeclaration }
 
             val newExpression = newCommonContainer.findExpressionByCopyableDataAndClearIt(EXPRESSION_KEY)
             val newCommonParent = newCommonContainer.findElementByCopyableDataAndClearIt(COMMON_PARENT_KEY)

idea/tests/org/jetbrains/kotlin/idea/ResolveElementCacheTest.kt

@@ -80,7 +80,7 @@ class C(param1: String = "", param2: Int = 0) {
         val klass = file.declarations.single() as KtClass
         val members = klass.declarations
         val function = members.first() as KtNamedFunction
-        val statements = (function.bodyExpression as KtBlockExpression).statements
+        val statements = function.bodyBlockExpression!!.statements
         return Data(file, klass, members, statements, KtPsiFactory(project))
     }
 
@@ -100,7 +100,7 @@ class C(param1: String = "", param2: Int = 0) {
         assert(aFunBodyContext3 === aFunBodyContext1)
 
         val bFun = members[1] as KtNamedFunction
-        val bBody = bFun.bodyExpression as KtBlockExpression
+        val bBody = bFun.bodyBlockExpression!!
         val bStatement = bBody.statements[0]
         val bFunBodyContext = bStatement.analyze(BodyResolveMode.FULL)
 
@@ -166,7 +166,7 @@ class C(param1: String = "", param2: Int = 0) {
 
             // modify body of "b()" via document
             val bFun = members[1] as KtNamedFunction
-            val bBody = bFun.bodyExpression as KtBlockExpression
+            val bBody = bFun.bodyBlockExpression!!
             document.insertString(bBody.lBrace!!.startOffset + 1, "x()")
             documentManager.commitAllDocuments()
 
@@ -427,7 +427,7 @@ class C(param1: String = "", param2: Int = 0) {
         val target = bindingContext[BindingContext.REFERENCE_TARGET, referenceExpr]
         TestCase.assertEquals("Ann", target?.importableFqName?.asString())
 
-        val statement = (function.bodyExpression as KtBlockExpression).statements[0]
+        val statement = function.bodyBlockExpression!!.statements[0]
         TestCase.assertEquals(null, bindingContext[BindingContext.PROCESSED, statement])
     }
 
@@ -465,17 +465,17 @@ class C(param1: String = "", param2: Int = 0) {
         val target = bindingContext[BindingContext.REFERENCE_TARGET, referenceExpr]
         TestCase.assertEquals("Ann", target?.importableFqName?.asString())
 
-        val statement = (constructor.bodyExpression as KtBlockExpression).statements[0]
+        val statement = constructor.bodyBlockExpression!!.statements[0]
         TestCase.assertEquals(null, bindingContext[BindingContext.PROCESSED, statement])
     }
 
     fun testFullResolveMultiple() {
         doTest {
-            val aBody = (members[0] as KtFunction).bodyExpression as KtBlockExpression
+            val aBody = (members[0] as KtFunction).bodyBlockExpression!!
             val statement1InFunA = aBody.statements[0]
             val statement2InFunA = aBody.statements[1]
-            val statementInFunB = ((members[1] as KtFunction).bodyExpression as KtBlockExpression).statements[0]
+            val statementInFunB = ((members[1] as KtFunction).bodyBlockExpression)!!.statements[0]
-            val statementInFunC = ((members[2] as KtFunction).bodyExpression as KtBlockExpression).statements[0]
+            val statementInFunC = ((members[2] as KtFunction).bodyBlockExpression)!!.statements[0]
 
             val bindingContext = checkResolveMultiple(BodyResolveMode.FULL, statement1InFunA, statementInFunB)
 
@@ -486,10 +486,10 @@ class C(param1: String = "", param2: Int = 0) {
 
     fun testPartialResolveMultiple() {
         doTest {
-            val aBody = (members[0] as KtFunction).bodyExpression as KtBlockExpression
+            val aBody = (members[0] as KtFunction).bodyBlockExpression!!
             val statement1InFunA = aBody.statements[0]
             val statement2InFunA = aBody.statements[1]
-            val statementInFunB = ((members[1] as KtFunction).bodyExpression as KtBlockExpression).statements[0]
+            val statementInFunB = ((members[1] as KtFunction).bodyBlockExpression)!!.statements[0]
             val constructorParameterDefault = klass.getPrimaryConstructor()!!.valueParameters[1].defaultValue!!
             val funC = members[2]
 
@@ -499,7 +499,7 @@ class C(param1: String = "", param2: Int = 0) {
 
     fun testPartialResolveMultipleInOneFunction() {
         doTest {
-            val aBody = (members[0] as KtFunction).bodyExpression as KtBlockExpression
+            val aBody = (members[0] as KtFunction).bodyBlockExpression!!
             val statement1InFunA = aBody.statements[0]
             val statement2InFunA = aBody.statements[1]
 

js/js.translator/src/org/jetbrains/kotlin/js/translate/utils/utils.kt

@@ -237,7 +237,7 @@ fun definePackageAlias(name: String, varName: JsName, tag: String, parentRef: Js
 val PsiElement.finalElement: PsiElement
     get() = when (this) {
         is KtFunctionLiteral -> rBrace ?: this
-        is KtDeclarationWithBody -> (bodyExpression as? KtBlockExpression)?.rBrace ?: bodyExpression ?: this
+        is KtDeclarationWithBody -> bodyBlockExpression?.rBrace ?: bodyExpression ?: this
         is KtLambdaExpression -> bodyExpression?.rBrace ?: this
         else -> this
     }