[FIR] Cleanup code in body resolve

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/BodyResolveContext.kt

@@ -70,9 +70,7 @@ class BodyResolveContext(
     @OptIn(PrivateForInline::class)
     inline fun <T> withNewTowerDataForClassParts(newContexts: FirTowerDataContextsForClassParts, f: () -> T): T {
         val old = towerDataContextsForClassParts
-
         towerDataContextsForClassParts = newContexts
-
         return try {
             f()
         } finally {
@@ -180,7 +178,7 @@ class BodyResolveContext(
     fun createSnapshotForLocalClasses(
         returnTypeCalculator: ReturnTypeCalculator,
         targetedLocalClasses: Set<FirClass<*>>
-    ) = BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
+    ): BodyResolveContext = BodyResolveContext(returnTypeCalculator, dataFlowAnalyzerContext, targetedLocalClasses, outerLocalClassForNested).apply {
         file = this@BodyResolveContext.file
         towerDataContextForAnonymousFunctions.putAll(this@BodyResolveContext.towerDataContextForAnonymousFunctions)
         containers = this@BodyResolveContext.containers

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirBodyResolveTransformer.kt

@@ -100,7 +100,7 @@ open class FirBodyResolveTransformer(
         wrappedArgumentExpression: FirWrappedArgumentExpression,
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
-        return (wrappedArgumentExpression.transformChildren(this, data) as FirStatement).compose()
+        return transformElement(wrappedArgumentExpression, data)
     }
 
     override fun transformQualifiedAccessExpression(
@@ -139,7 +139,10 @@ open class FirBodyResolveTransformer(
         return expressionsTransformer.transformComparisonExpression(comparisonExpression, data)
     }
 
-    override fun transformTypeOperatorCall(typeOperatorCall: FirTypeOperatorCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformTypeOperatorCall(
+        typeOperatorCall: FirTypeOperatorCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return expressionsTransformer.transformTypeOperatorCall(typeOperatorCall, data)
     }
 
@@ -189,7 +192,10 @@ open class FirBodyResolveTransformer(
         return declarationsTransformer.transformWrappedDelegateExpression(wrappedDelegateExpression, data)
     }
 
-    override fun <T> transformConstExpression(constExpression: FirConstExpression<T>, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun <T> transformConstExpression(
+        constExpression: FirConstExpression<T>,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return expressionsTransformer.transformConstExpression(constExpression, data)
     }
 
@@ -204,7 +210,10 @@ open class FirBodyResolveTransformer(
         return expressionsTransformer.transformDelegatedConstructorCall(delegatedConstructorCall, data)
     }
 
-    override fun transformAugmentedArraySetCall(augmentedArraySetCall: FirAugmentedArraySetCall, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformAugmentedArraySetCall(
+        augmentedArraySetCall: FirAugmentedArraySetCall,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return expressionsTransformer.transformAugmentedArraySetCall(augmentedArraySetCall, data)
     }
 
@@ -249,15 +258,24 @@ open class FirBodyResolveTransformer(
         return declarationsTransformer.transformRegularClass(regularClass, data)
     }
 
-    override fun transformAnonymousObject(anonymousObject: FirAnonymousObject, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformAnonymousObject(
+        anonymousObject: FirAnonymousObject,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return declarationsTransformer.transformAnonymousObject(anonymousObject, data)
     }
 
-    override fun transformSimpleFunction(simpleFunction: FirSimpleFunction, data: ResolutionMode): CompositeTransformResult<FirSimpleFunction> {
+    override fun transformSimpleFunction(
+        simpleFunction: FirSimpleFunction,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirSimpleFunction> {
         return declarationsTransformer.transformSimpleFunction(simpleFunction, data)
     }
 
-    override fun <F : FirFunction<F>> transformFunction(function: FirFunction<F>, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun <F : FirFunction<F>> transformFunction(
+        function: FirFunction<F>,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return declarationsTransformer.transformFunction(function, data)
     }
 
@@ -272,7 +290,10 @@ open class FirBodyResolveTransformer(
         return declarationsTransformer.transformAnonymousInitializer(anonymousInitializer, data)
     }
 
-    override fun transformAnonymousFunction(anonymousFunction: FirAnonymousFunction, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformAnonymousFunction(
+        anonymousFunction: FirAnonymousFunction,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return declarationsTransformer.transformAnonymousFunction(anonymousFunction, data)
     }
 
@@ -317,11 +338,17 @@ open class FirBodyResolveTransformer(
         return controlFlowStatementsTransformer.transformJump(jump, data)
     }
 
-    override fun transformThrowExpression(throwExpression: FirThrowExpression, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformThrowExpression(
+        throwExpression: FirThrowExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return controlFlowStatementsTransformer.transformThrowExpression(throwExpression, data)
     }
 
-    override fun transformElvisExpression(elvisExpression: FirElvisExpression, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformElvisExpression(
+        elvisExpression: FirElvisExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         return controlFlowStatementsTransformer.transformElvisExpression(elvisExpression, data)
     }
 

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirControlFlowStatementsResolveTransformer.kt

@@ -45,9 +45,8 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
         // Do-while has a specific scope structure (its block and condition effectively share the scope)
         return withNewLocalScope {
             doWhileLoop.also(dataFlowAnalyzer::enterDoWhileLoop)
-                .apply {
+                .also {
-                    transformer.expressionsTransformer.transformBlockInCurrentScope(block, context)
+                    transformer.expressionsTransformer.transformBlockInCurrentScope(it.block, context)
-                    dataFlowAnalyzer
                 }
                 .also(dataFlowAnalyzer::enterDoWhileLoopCondition).transformCondition(transformer, context)
                 .also(dataFlowAnalyzer::exitDoWhileLoop)
@@ -140,7 +139,7 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
             return tryExpression.compose()
         }
 
-        tryExpression.annotations.forEach { it.accept(this, data) }
+        tryExpression.transformAnnotations(transformer, ResolutionMode.ContextIndependent)
         dataFlowAnalyzer.enterTryExpression(tryExpression)
         tryExpression.transformTryBlock(transformer, ResolutionMode.ContextDependent)
         dataFlowAnalyzer.exitTryMainBlock(tryExpression)
@@ -211,7 +210,10 @@ class FirControlFlowStatementsResolveTransformer(transformer: FirBodyResolveTran
 
     // ------------------------------- Elvis -------------------------------
 
-    override fun transformElvisExpression(elvisExpression: FirElvisExpression, data: ResolutionMode): CompositeTransformResult<FirStatement> {
+    override fun transformElvisExpression(
+        elvisExpression: FirElvisExpression,
+        data: ResolutionMode
+    ): CompositeTransformResult<FirStatement> {
         if (elvisExpression.calleeReference is FirResolvedNamedReference) return elvisExpression.compose()
         elvisExpression.transformAnnotations(transformer, data)
         elvisExpression.transformLhs(transformer, ResolutionMode.ContextDependent)

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirDeclarationsResolveTransformer.kt

@@ -410,7 +410,7 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
             if (anonymousObject.classKind == ClassKind.ENUM_ENTRY) {
                 anonymousObject.getPrimaryConstructorIfAny()?.symbol?.callableId?.className?.shortName()
             } else null
-        var result = withScopesForClass(labelName, anonymousObject, type) {
+        val result = withScopesForClass(labelName, anonymousObject, type) {
             transformDeclarationContent(anonymousObject, data).single as FirAnonymousObject
         }
         if (!implicitTypeOnly && result.controlFlowGraphReference == null) {
@@ -431,12 +431,10 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                 lambdaResolution.expectedReturnTypeRef ?: anonymousFunction.returnTypeRef.takeUnless { it is FirImplicitTypeRef }
             val result = transformFunction(anonymousFunction, withExpectedType(expectedReturnType)).single as FirAnonymousFunction
             val body = result.body
-            return if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
+            if (result.returnTypeRef is FirImplicitTypeRef && body != null) {
                 result.transformReturnTypeRef(transformer, withExpectedType(body.resultType))
-                result
-            } else {
-                result
             }
+            return result
         }
 
         val label = anonymousFunction.label
@@ -492,7 +490,6 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
         function: F,
         resolutionMode: ResolutionMode,
     ): CompositeTransformResult<F> {
-
         @Suppress("UNCHECKED_CAST")
         val result = transformFunction(function, resolutionMode).single as F
 
@@ -691,15 +688,15 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                         it.type, it, anonymousFunction, returnTypeVariable = null, components, candidate = null
                     )
                 }
-                var af = anonymousFunction
+                var lambda = anonymousFunction
-                val valueParameters =
+                val valueParameters = when (resolvedLambdaAtom) {
-                    if (resolvedLambdaAtom == null) af.valueParameters
+                    null -> lambda.valueParameters
-                    else {
+                    else -> {
                         val singleParameterType = resolvedLambdaAtom.parameters.singleOrNull()
-                        val itParam = when {
+                        when {
-                            af.valueParameters.isEmpty() && singleParameterType != null -> {
+                            lambda.valueParameters.isEmpty() && singleParameterType != null -> {
                                 val name = Name.identifier("it")
-                                buildValueParameter {
+                                val itParam = buildValueParameter {
                                     session = this@FirDeclarationsResolveTransformer.session
                                     origin = FirDeclarationOrigin.Source
                                     returnTypeRef = buildResolvedTypeRef { type = singleParameterType }
@@ -709,42 +706,42 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                                     isNoinline = false
                                     isVararg = false
                                 }
+                                listOf(itParam)
                             }
-                            else -> null
+
-                        }
+                            else -> {
-                        if (itParam != null) {
+                                lambda.valueParameters.mapIndexed { index, param ->
-                            listOf(itParam)
+                                    if (param.returnTypeRef is FirResolvedTypeRef) {
-                        } else {
+                                        param
-                            af.valueParameters.mapIndexed { index, param ->
+                                    } else {
-                                if (param.returnTypeRef is FirResolvedTypeRef) {
+                                        param.transformReturnTypeRef(
-                                    param
+                                            StoreType,
-                                } else {
+                                            param.returnTypeRef.resolvedTypeFromPrototype(
-                                    param.transformReturnTypeRef(
+                                                resolvedLambdaAtom.parameters[index]
-                                        StoreType,
+                                            )
-                                        param.returnTypeRef.resolvedTypeFromPrototype(
-                                            resolvedLambdaAtom.parameters[index]
                                         )
-                                    )
+                                        param
-                                    param
+                                    }
                                 }
                             }
                         }
-
                     }
-                val returnTypeRefFromResolvedAtom = resolvedLambdaAtom?.returnType?.let { af.returnTypeRef.resolvedTypeFromPrototype(it) }
+                }
-                af = af.copy(
+                val returnTypeRefFromResolvedAtom =
-                    receiverTypeRef = af.receiverTypeRef?.takeIf { it !is FirImplicitTypeRef }
+                    resolvedLambdaAtom?.returnType?.let { lambda.returnTypeRef.resolvedTypeFromPrototype(it) }
-                        ?: resolvedLambdaAtom?.receiver?.let { af.receiverTypeRef?.resolvedTypeFromPrototype(it) },
+                lambda = lambda.copy(
+                    receiverTypeRef = lambda.receiverTypeRef?.takeIf { it !is FirImplicitTypeRef }
+                        ?: resolvedLambdaAtom?.receiver?.let { lambda.receiverTypeRef?.resolvedTypeFromPrototype(it) },
                     valueParameters = valueParameters,
-                    returnTypeRef = (af.returnTypeRef as? FirResolvedTypeRef)
+                    returnTypeRef = (lambda.returnTypeRef as? FirResolvedTypeRef)
                         ?: returnTypeRefFromResolvedAtom
-                        ?: af.returnTypeRef
+                        ?: lambda.returnTypeRef
                 )
-                af = af.transformValueParameters(ImplicitToErrorTypeTransformer, null)
+                lambda = lambda.transformValueParameters(ImplicitToErrorTypeTransformer, null)
                 val bodyExpectedType = returnTypeRefFromResolvedAtom ?: expectedTypeRef
-                val labelName = af.label?.name?.let { Name.identifier(it) }
+                val labelName = lambda.label?.name?.let { Name.identifier(it) }
-                withLabelAndReceiverType(labelName, af, af.receiverTypeRef?.coneType) {
+                withLabelAndReceiverType(labelName, lambda, lambda.receiverTypeRef?.coneType) {
-                    af = transformFunction(af, withExpectedType(bodyExpectedType)).single as FirAnonymousFunction
+                    lambda = transformFunction(lambda, withExpectedType(bodyExpectedType)).single as FirAnonymousFunction
                 }
                 // To separate function and separate commit
                 val writer = FirCallCompletionResultsWriterTransformer(
@@ -755,20 +752,20 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
                     integerOperatorsTypeUpdater,
                     integerLiteralTypeApproximator
                 )
-                af.transformSingle(writer, expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.toExpectedType())
+                lambda.transformSingle(writer, expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.toExpectedType())
-                val returnTypes = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(af)
+                val returnTypes = dataFlowAnalyzer.returnExpressionsOfAnonymousFunction(lambda)
                     .mapNotNull { (it as? FirExpression)?.resultType?.coneType }
-                af.replaceReturnTypeRef(
+                lambda.replaceReturnTypeRef(
-                    af.returnTypeRef.resolvedTypeFromPrototype(
+                    lambda.returnTypeRef.resolvedTypeFromPrototype(
                         inferenceComponents.ctx.commonSuperTypeOrNull(returnTypes) ?: session.builtinTypes.unitType.type
                     )
                 )
-                af.replaceTypeRef(
+                lambda.replaceTypeRef(
-                    af.constructFunctionalTypeRef(
+                    lambda.constructFunctionalTypeRef(
                         isSuspend = expectedTypeRef.coneTypeSafe<ConeKotlinType>()?.isSuspendFunctionType(session) == true
                     )
                 )
-                af.addReturn().compose()
+                lambda.addReturn().compose()
             }
             is ResolutionMode.WithStatus -> {
                 throw AssertionError("Should not be here in WithStatus mode")
@@ -861,12 +858,11 @@ open class FirDeclarationsResolveTransformer(transformer: FirBodyResolveTransfor
 
         components.context.replaceTowerDataContext(forMembersResolution)
 
-        val newContexts =
+        val newContexts = FirTowerDataContextsForClassParts(
-            FirTowerDataContextsForClassParts(
+            newTowerDataContextForStaticNestedClasses,
-                newTowerDataContextForStaticNestedClasses,
+            scopeForConstructorHeader,
-                scopeForConstructorHeader,
+            primaryConstructorParametersScope
-                primaryConstructorParametersScope
+        )
-            )
 
         context.withNewTowerDataForClassParts(newContexts) {
             block()

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/body/resolve/FirExpressionsResolveTransformer.kt

@@ -50,8 +50,10 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         if (expression.resultType is FirImplicitTypeRef && expression !is FirWrappedExpression) {
             val type = buildErrorTypeRef {
                 source = expression.source
-                diagnostic =
+                diagnostic = ConeSimpleDiagnostic(
-                    ConeSimpleDiagnostic("Type calculating for ${expression::class} is not supported", DiagnosticKind.InferenceError)
+                    "Type calculating for ${expression::class} is not supported",
+                    DiagnosticKind.InferenceError
+                )
             }
             expression.resultType = type
         }
@@ -260,8 +262,8 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         if (functionCall.calleeReference !is FirSimpleNamedReference) return functionCall.compose()
         if (functionCall.calleeReference is FirNamedReferenceWithCandidate) return functionCall.compose()
         dataFlowAnalyzer.enterCall()
-        functionCall.annotations.forEach { it.accept(this, data) }
+        functionCall.transformAnnotations(transformer, data)
-        functionCall.transform<FirFunctionCall, Nothing?>(InvocationKindTransformer, null)
+        functionCall.transformSingle(InvocationKindTransformer, null)
         functionCall.transformTypeArguments(transformer, ResolutionMode.ContextIndependent)
         val expectedTypeRef = data.expectedType
         val (completeInference, callCompleted) =
@@ -302,9 +304,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
 
         block.transformStatementsIndexed(transformer) { index ->
             val value = if (index == numberOfStatements - 1) data else ResolutionMode.ContextIndependent
-
             transformer.onBeforeStatementResolution(block.statements[index])
-
             TransformData.Data(value)
         }
         if (data == ResolutionMode.ContextIndependent) {
@@ -334,7 +334,8 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
     ): CompositeTransformResult<FirStatement> {
         return (comparisonExpression.transformChildren(transformer, ResolutionMode.ContextIndependent) as FirComparisonExpression).also {
             it.resultType = comparisonExpression.typeRef.resolvedTypeFromPrototype(builtinTypes.booleanType.type)
-        }.transformSingle(integerLiteralTypeApproximator, null).also(dataFlowAnalyzer::exitComparisonExpressionCall).compose()
+        }.transformSingle(integerLiteralTypeApproximator, null)
+            .also(dataFlowAnalyzer::exitComparisonExpressionCall).compose()
     }
 
     override fun transformAssignmentOperatorStatement(
@@ -347,7 +348,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         val leftArgument = assignmentOperatorStatement.leftArgument.transformSingle(transformer, ResolutionMode.ContextIndependent)
         val rightArgument = assignmentOperatorStatement.rightArgument.transformSingle(transformer, ResolutionMode.ContextDependent)
 
-        fun createFunctionCall(name: Name) = buildFunctionCall {
+        fun createFunctionCall(name: Name): FirFunctionCall = buildFunctionCall {
             source = assignmentOperatorStatement.source?.fakeElement(FirFakeSourceElementKind.DesugaredCompoundAssignment)
             explicitReceiver = leftArgument
             argumentList = buildUnaryArgumentList(rightArgument)
@@ -424,7 +425,6 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         equalityOperatorCall: FirEqualityOperatorCall,
         data: ResolutionMode
     ): CompositeTransformResult<FirStatement> {
-        // TODO: add approximation of integer literals
         val result = (equalityOperatorCall.transformChildren(transformer, ResolutionMode.ContextIndependent) as FirEqualityOperatorCall)
             .also { it.resultType = equalityOperatorCall.typeRef.resolvedTypeFromPrototype(builtinTypes.booleanType.type) }
             .transformSingle(integerLiteralTypeApproximator, null)
@@ -451,7 +451,8 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         val baseTypeArguments =
             argument.typeRef.coneTypeSafe<ConeKotlinType>()?.fullyExpandedType(session)?.typeArguments
         val type = coneTypeSafe<ConeKotlinType>()
-        return if (type?.typeArguments?.isEmpty() != true ||
+        return if (
+            type?.typeArguments?.isEmpty() != true ||
             type is ConeTypeParameterType ||
             baseTypeArguments?.isEmpty() != false ||
             (type is ConeClassLikeType &&
@@ -471,6 +472,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             typeOperatorCall.transformConversionTypeRef(transformer, ResolutionMode.ContextIndependent)
         }.transformOtherChildren(transformer, ResolutionMode.ContextIndependent)
         resolved.argumentList.transformArguments(integerLiteralTypeApproximator, null)
+
         val conversionTypeRef = resolved.conversionTypeRef.withTypeArgumentsForBareType(resolved.argument)
         resolved.transformChildren(object : FirDefaultTransformer<Nothing?>() {
             override fun <E : FirElement> transformElement(element: E, data: Nothing?): CompositeTransformResult<E> {
@@ -485,6 +487,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                 }
             }
         }, null)
+
         when (resolved.operation) {
             FirOperation.IS, FirOperation.NOT_IS -> {
                 resolved.resultType = session.builtinTypes.booleanType
@@ -504,7 +507,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
                         ),
                     )
             }
-            else -> error("Unknown type operator")
+            else -> error("Unknown type operator: ${resolved.operation}")
         }
         dataFlowAnalyzer.exitTypeOperatorCall(resolved)
         return resolved.transform(integerLiteralTypeApproximator, null)
@@ -591,7 +594,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
             return callableReferenceAccess.compose()
         }
 
-        callableReferenceAccess.annotations.forEach { it.accept(this, data) }
+        callableReferenceAccess.transformAnnotations(transformer, data)
         val explicitReceiver = callableReferenceAccess.explicitReceiver
         val transformedLHS = explicitReceiver?.transformSingle(this, ResolutionMode.ContextIndependent)?.apply {
             if (this is FirResolvedQualifier && callableReferenceAccess.hasQuestionMarkAtLHS) {
@@ -872,6 +875,7 @@ open class FirExpressionsResolveTransformer(transformer: FirBodyResolveTransform
         assert(augmentedArraySetCall.operation in FirOperation.ASSIGNMENTS)
         assert(augmentedArraySetCall.operation != FirOperation.ASSIGN)
 
+        augmentedArraySetCall.transformAnnotations(transformer, data)
         val operatorName = FirOperationNameConventions.ASSIGNMENTS.getValue(augmentedArraySetCall.operation)
 
         val firstCalls = with(augmentedArraySetCall.setGetBlock.statements.last() as FirFunctionCall) setCall@{