Introduce helper FirExpression.unwrapArgument() in fir:tree

compiler/fir/checkers/src/org/jetbrains/kotlin/fir/analysis/checkers/expression/FirAnnotationUsedAsAnnotationArgumentChecker.kt

@@ -10,13 +10,13 @@ import org.jetbrains.kotlin.fir.analysis.diagnostics.DiagnosticReporter
 import org.jetbrains.kotlin.fir.analysis.diagnostics.FirErrors
 import org.jetbrains.kotlin.fir.analysis.diagnostics.reportOn
 import org.jetbrains.kotlin.fir.expressions.FirAnnotationCall
-import org.jetbrains.kotlin.fir.expressions.FirWrappedArgumentExpression
+import org.jetbrains.kotlin.fir.expressions.unwrapArgument
 
 object FirAnnotationUsedAsAnnotationArgumentChecker : FirAnnotationCallChecker() {
     override fun check(expression: FirAnnotationCall, context: CheckerContext, reporter: DiagnosticReporter) {
         val args = expression.argumentList.arguments
         for (arg in args) {
-            for (ann in ((arg as? FirWrappedArgumentExpression)?.expression ?: arg).annotations) {
+            for (ann in arg.unwrapArgument().annotations) {
                 reporter.reportOn(ann.source, FirErrors.ANNOTATION_USED_AS_ANNOTATION_ARGUMENT, context)
             }
         }

compiler/fir/fir2ir/src/org/jetbrains/kotlin/fir/backend/Fir2IrVisitor.kt

@@ -459,11 +459,7 @@ class Fir2IrVisitor(
                 )
             }
             else -> {
-                val unwrappedExpression = if (expression is FirWrappedArgumentExpression) {
-                    expression.expression
-                } else {
-                    expression
-                }
+                val unwrappedExpression = expression.unwrapArgument()
                 if (annotationMode && unwrappedExpression is FirFunctionCall) {
                     convertToIrCall(unwrappedExpression, annotationMode)
                 } else {

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/calls/Arguments.kt

@@ -494,7 +494,7 @@ fun FirExpression.isFunctional(
     scopeSession: ScopeSession,
     expectedFunctionType: ConeKotlinType?,
 ): Boolean {
-    when ((this as? FirWrappedArgumentExpression)?.expression ?: this) {
+    when (unwrapArgument()) {
         is FirAnonymousFunction, is FirCallableReferenceAccess -> return true
         else -> {
             // Either a functional type or a subtype of a class that has a contributed `invoke`.

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/contracts/ConeConditionalEffectToFirVisitor.kt

@@ -103,7 +103,7 @@ fun createArgumentsMapping(qualifiedAccess: FirQualifiedAccess): Map<Int, FirExp
             val parameterToIndex = function.valueParameters.mapIndexed { index, parameter -> parameter to index }.toMap()
             val callArgumentMapping = qualifiedAccess.argumentMapping ?: return null
             for (argument in qualifiedAccess.arguments) {
-                argumentsMapping[parameterToIndex.getValue(callArgumentMapping.getValue(argument))] = argument.unwrap()
+                argumentsMapping[parameterToIndex.getValue(callArgumentMapping.getValue(argument))] = argument.unwrapArgument()
             }
         }
         is FirVariableAssignment -> {
@@ -113,4 +113,3 @@ fun createArgumentsMapping(qualifiedAccess: FirQualifiedAccess): Map<Int, FirExp
     return argumentsMapping
 }
 
-private fun FirExpression.unwrap(): FirExpression = if (this is FirWrappedArgumentExpression) expression else this

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirCallCompletionResultsWriterTransformer.kt

@@ -755,11 +755,6 @@ private fun ExpectedArgumentType.getExpectedType(argument: FirExpression): ConeK
 
 fun ConeKotlinType.toExpectedType(): ExpectedArgumentType = ExpectedArgumentType.ExpectedType(this)
 
-private fun FirExpression.unwrapArgument(): FirExpression = when (this) {
-    is FirWrappedArgumentExpression -> expression
-    else -> this
-}
-
 class FirDeclarationCompletionResultsWriter(private val finalSubstitutor: ConeSubstitutor) : FirDefaultTransformer<Any?>() {
     override fun <E : FirElement> transformElement(element: E, data: Any?): E {
         return element

compiler/fir/tree/src/org/jetbrains/kotlin/fir/expressions/FirExpressionUtil.kt

@@ -108,3 +108,5 @@ fun FirBlock.replaceFirstStatement(statement: FirStatement): FirStatement {
     statements[0] = statement
     return existed
 }
+
+fun FirExpression.unwrapArgument(): FirExpression = (this as? FirWrappedArgumentExpression)?.expression ?: this